public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OptimizeBatchMatmulPass)

  OptimizeBatchMatmulPass() = default;
  OptimizeBatchMatmulPass(const OptimizeBatchMatmulPass &) {}

  void runOnOperation() override;
};

// Converts batch_matmul operation to fully_connected if rhs is a
// constant tensor with rank 2
struct ConvertBatchMatMulOp2FullyConnectedOp
    : public OpRewritePattern<TFL::BatchMatMulOp> {

// pass and folders are not applied by default.
class UnfoldSplatConstantPass
    : public impl::UnfoldSplatConstantPassBase<UnfoldSplatConstantPass> {
 public:
  void runOnOperation() override {
    auto module = getOperation();

    mlir::OpBuilder op_builder(&module.getBodyRegion());
    // Cannot use the pattern rewriter because the driver applies folders by
    // default.

// the end.
struct UnrollBatchMatMulPass
    : public impl::UnrollBatchMatMulPassBase<UnrollBatchMatMulPass> {
  void runOnOperation() override;
};

void UnrollBatchMatMulPass::runOnOperation() {
  RewritePatternSet patterns(&getContext());
  auto func = getOperation();
  PopulateUnrollTfBatchMatMul(&getContext(), patterns);

    enable_weight_only_quantization_ = quant_specs.weight_only_quantization;
  }

  void runOnOperation() override;

 private:
  quant::QuantizationSpecs quant_specs;
};

#include "tensorflow/compiler/mlir/lite/transforms/generated_quantize.inc"

void QuantizePass::runOnOperation() {
  RewritePatternSet patterns(&getContext());
  auto func = getOperation();
  auto* ctx = func.getContext();

  }

  StringRef getDescription() const override {
    return "Merge the save function's ops to the main function. The save "
           "function will be removed after the pass.";
  }

  void runOnOperation() override;
};

// Returns true iff func_op has either no Region or the body has no Blocks.
bool IsFuncOpEmpty(func::FuncOp func_op) {
  return func_op->getNumRegions() == 0 || func_op.getBody().empty();
}

#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

struct TPUDevicePropagation
    : public impl::TPUDevicePropagationPassBase<TPUDevicePropagation> {
  void runOnOperation() override;
};

void TPUDevicePropagation::runOnOperation() {
  func::FuncOp func = getOperation();
  if (!IsSupportedGraph(func)) return;

  llvm::DenseMap<Value, llvm::StringRef> value_to_device;

    return "quant-preprocess-op";
  }
  StringRef getDescription() const final {
    // This is a brief description of the pass.
    return "Preprocess TF op prior to quantization";
  }

  void runOnOperation() override;

 private:
  Option<OpSet> op_set_{
      *this, "target-opset", llvm::cl::init(OpSet::UNIFORM_QUANTIZED),
      llvm::cl::desc("Choose target opset."),
      llvm::cl::values(

    if (tf2xla_fallback_device_type.has_value()) {
      device_type_ = tf2xla_fallback_device_type.value().str();
    }
  }
  /// Performs the lowering to XLA dialect.
  void runOnOperation() override;
};

#define GEN_PASS_DEF_LEGALIZETFMODULEPASS
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_tf_passes.h.inc"

}  // namespace

class DeferActivationTransposePass
    : public impl::DeferActivationTransposePassBase<
          DeferActivationTransposePass> {
 private:
  void runOnOperation() override;
};

void DeferActivationTransposePass::runOnOperation() {
  func::FuncOp func_op = getOperation();
  MLIRContext& ctx = getContext();

  RewritePatternSet patterns(&ctx);
  patterns.add<DeferActivationTransposeForAddOp,

  // processing the function for dynamic range quantization. Therefore, this
  // method preprocess the function to remove all stats ops.
  void removeAllStatsOp(func::FuncOp func);

  void runOnOperation() override;

 private:
  QuantizationSpecs quant_specs_;
  OpSet op_set_;

  Option<bool> enable_per_channel_quantization_{
      *this, "enable-per-channel-quantization", llvm::cl::init(false),