};

// Unroll split into a bunch of slice ops.
struct UnrollSplit : public OpRewritePattern<TFL::SplitOp> {
  using OpRewritePattern<TFL::SplitOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(TFL::SplitOp split_op,
                                PatternRewriter& rewriter) const override;
};

// Unroll splitv into a bunch of slice ops.

      mlir::TF::SplitOp root_split_op;
      auto result = CreateSplitOp(num_splits, dimension, location,
                                  original_source, builder, &root_split_op);
      if (mlir::failed(result)) return mlir::failure();

      split_ops_for_tiled_input.emplace_back(root_split_op);
      continue;
    }

    llvm::SmallVector<mlir::TF::SplitOp, 4> new_split_ops;

  void runOnOperation() override;
};

void TPUColocateSplits::runOnOperation() {
  getOperation().walk([&](Operation* op) {
    if (auto split = llvm::dyn_cast<TF::SplitOp>(op)) {
      if (HasDevice(split) || split->getAttrOfType<ArrayAttr>(kClassAttr))
        return WalkResult::advance();
      for (Operation* pred : IslandPredecessors(split)) {
        if (auto colocation_classes =

bool NotFromDequant(mlir::Value value) {
  auto dequant_op = value.getDefiningOp<DequantizeOp>();
  if (dequant_op) {
    return false;
  }
  auto split_op = value.getDefiningOp<SplitOp>();
  if (!split_op) {
    return true;
  }
  return !split_op.getValue().getDefiningOp<DequantizeOp>();
}

// Optimize TFLite operations in functions.
class OptimizeBatchMatmulPass

    Type slice_result_type = RankedTensorType::get(slice_size, element_type);
    llvm::SmallVector<Type, 4> output_types(batch_size, slice_result_type);
    auto split_op = rewriter.create<TF::SplitOp>(loc, output_types,
                                                 split_dimension_op.getOutput(),
                                                 reshape_op.getOutput());

    // Squeeze each batch, i.e. reshape

                           ConcatV2Op::getOperationName(),
                           AddV2Op::getOperationName(),
                           PadOp::getOperationName(),
                           SplitOp::getOperationName(),
                           UnpackOp::getOperationName(),
                           DivOp::getOperationName(),
                           MulOp::getOperationName(),

      CompileTimeConstantOperand<TF::SliceOp, 1, 2>,       // $begin, $size
      CompileTimeConstantOperand<TF::SparseToDenseOp, 1>,  // $output_shape
      CompileTimeConstantOperand<TF::SplitOp, 0>,          // $split_dim
      // $size_splits, $split_dim
      CompileTimeConstantOperand<TF::SplitVOp, 1, 2>,
      CompileTimeConstantOperand<TF::StackV2Op, 0>,  // $max_size
      // $num_samples

LogicalResult ConvertTFSplitOp::matchAndRewrite(
    Operation* op, PatternRewriter& rewriter) const {
  auto tf_split_op = cast<TF::SplitOp>(op);

  // Number of splits cannot be negative.
  auto num_split = rewriter.getI32IntegerAttr(tf_split_op.getNumSplit());

  rewriter.replaceOpWithNewOp<TFL::SplitOp>(
      op, tf_split_op.getOutput().getTypes(), tf_split_op.getSplitDim(),
      tf_split_op.getValue(), num_split);

    TypeID::get<TF::SoftmaxCrossEntropyWithLogitsOp>(),
    TypeID::get<TF::SoftplusOp>(),
    TypeID::get<TF::SparseMatMulOp>(),
    TypeID::get<TF::SparseSoftmaxCrossEntropyWithLogitsOp>(),
    TypeID::get<TF::SplitOp>(),
    TypeID::get<TF::SplitVOp>(),
    TypeID::get<TF::SqrtGradOp>(),
    TypeID::get<TF::SquaredDifferenceOp>(),
    TypeID::get<TF::SqueezeOp>(),
    TypeID::get<TF::StatelessParameterizedTruncatedNormalOp>(),

  rewriter.replaceOp(squared_diff_op, mul_op.getResult());

  return success();
}

// ================== split ========================

LogicalResult UnrollSplit::matchAndRewrite(TFL::SplitOp split_op,
                                           PatternRewriter& rewriter) const {
  auto num_splits = split_op.getNumSplits();
  auto input = split_op.getValue();