return %0 : tensor<2x2xf32>
    }
    func.func private @composite_fn_1(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> attributes {tf_quant.composite_function} {

    bool enable_canonicalization, bool disable_fuse_mul_and_fc = false);
std::unique_ptr<OperationPass<func::FuncOp>> CreateOptimizePass();

// Creates an instance of the Tensorflow Lite batch matmul Optimize pass.
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", "false",
             "Disable folding mul and fully connected ops during optimization pass.">,
  ];
}

def OptimizeBatchMatmulPass : Pass<"tfl-optimize-batch-matmul", "mlir::func::FuncOp"> {
  let summary = "Optimize FC with BatchMatmul within the TensorFlow Lite dialect";
  let constructor = "CreateOptimizeBatchMatmulPass()";
  let dependentDialects = ["TFL::TensorFlowLiteDialect"];
}

    // 262657 = 1 + 512 + 512 * 512
    // expected-remark@+1 {{Cost: 262657}}
    %2 = "tf.MatMul"(%arg, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<?x512xf32>, tensor<512x512xf32>) -> tensor<?x512xf32>
    // expected-remark@+1 {{Cost: 512}}

  %7 = "tf.Reshape"(%6, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<9xf32>, tensor<2xi32>) -> tensor<3x3xf32>
  %8 = "tf.MatMul"(%2, %7) {device = "/job:localhost/replica:0/task:0/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>

        TypeID::get<TF::LogicalNotOp>(),
        TypeID::get<TF::LogOp>(),
        TypeID::get<TF::LowerBoundOp>(),
        TypeID::get<TF::MakeUniqueOp>(),
        TypeID::get<TF::MatMulOp>(),
        TypeID::get<TF::MatrixDiagV3Op>(),
        TypeID::get<TF::MatrixInverseOp>(),
        TypeID::get<TF::MatrixSetDiagV3Op>(),
        TypeID::get<TF::MatrixSolveOp>(),

  %x = "tf.TensorArrayReadV3"(%handle_0, %index, %flow_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<i32>, tensor<f32>) -> tensor<?x100xf32>
  %y = "tf.MatMul"(%x, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x100xf32>, tensor<100x512xf32>) -> (tensor<?x512xf32>)

// identifier is the order of that attribute in `attributes`. This map
// is then used to set attributes in the quantized functions in the
// QuantizeCompositeFunctionsPass.
// For example, for tf.MatMul with `attributes` = {{"transpose_a", false},
// {"transpose_b", false}}, the generated attr_map is
// "0:transpose_a,1:transpose_b", where 0 and 1 are the respective attribute
// identifiers.

          "(LRN, LRNGrad)."
          " BN: TF FusedBatchNorm* operations."
          " FUSIBLE: All TF operations that XLA can fuse (All the above). "
          "You can also put any TF operation name, e.g. 'FUSIBLE,MatMul'."),
      Flag("tf_xla_cluster_exclude_ops",
           &mark_for_compilation_flags->tf_xla_cluster_exclude_ops,
           "(experimental) "
           "Exclude the operations from auto-clustering. "

    %16 = "tf.Mean"(%15, %8) {keep_dims = false} : (tensor<2x112x112x64xf32>, tensor<2xi32>) -> tensor<2x64xf32>
    %17 = "tf.MatMul"(%16, %arg3) {transpose_a = false, transpose_b = false} : (tensor<2x64xf32>, tensor<64x1001xf32>) -> tensor<2x1001xf32>
    %18 = "tf.BiasAdd"(%17, %arg4) {data_format = "NHWC"} : (tensor<2x1001xf32>, tensor<1001xf32>) -> tensor<2x1001xf32>