};

// 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.

  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