mlir::TFL::CreateOptimizeOpOrderPass());
  // Add optimization pass after quantization for additional fusing
  // opportunities.

  if (!pass_config.unfold_batch_matmul) {
    // Enable an optimization pass that transforms FC to BatchMatmul only when
    // `unfold_batch_matmul=false`.
    pass_manager.addNestedPass<mlir::func::FuncOp>(
        mlir::TFL::CreateOptimizeBatchMatmulPass());
  }

  PrepareTFPass() = default;
  PrepareTFPass(const PrepareTFPass &) {}
  explicit PrepareTFPass(bool unfold_batch_matmul,
                         bool allow_bf16_and_f16_type_legalization,
                         bool use_fake_quant_num_bits = false) {
    this->unfold_batch_matmul_ = unfold_batch_matmul;
    this->allow_bf16_and_f16_type_legalization_ =
        allow_bf16_and_f16_type_legalization;

    // function.
    // And in fact, we probably need to do the proper legalization, for the
    // compute cost to work. (in case we added some TF ops)
    pass_manager->addPass(mlir::TFL::CreatePrepareTFPass(
        /*unfold_batch_matmul=*/true,
        /*allow_bf16_and_f16_type_legalization=*/false));
    pass_manager->addNestedPass<mlir::func::FuncOp>(
        mlir::createCanonicalizerPass());
    pass_manager->addPass(

std::unique_ptr<OperationPass<func::FuncOp>> CreateOptimizeBatchMatmulPass();

// Creates an instance of the TensorFlow Lite dialect PrepareTF pass.
std::unique_ptr<OperationPass<func::FuncOp>> CreatePrepareTFPass(
    bool unfold_batch_matmul, bool allow_bf16_and_f16_type_legalization,
    bool use_fake_quant_num_bits = false);
std::unique_ptr<OperationPass<func::FuncOp>> CreatePrepareTFPass();