def MakeOneDimValueBroadcastable : NativeCodeCall<
  "MakeOneDimValueBroadcastable($_builder, $_loc, $0, $1.getType().cast<ShapedType>())">;

// Match convolution op with "NHWC" data format or matmul op.
def SupportedAffineOpMatcher : NativeCodeCall<
  "MatchSupportedAffineOp($_self, $0, $1, $2)">;

// Checks if a value can be symetrically quantized.

#include "xla/xla_data.pb.h"

namespace mlir::quant {
namespace {

constexpr StringRef kTfQuantCreatedEinsum = "__tf_quant_created_einsum";

// Replaces mixed-type Conv and Matmul cast hacks with TF XLA ops.
// TODO(b/228403741): Support conversion for dynamic-shaped TF ops.
class ReplaceCastHacksWithTFXLAOpsPass
    : public PassWrapper<ReplaceCastHacksWithTFXLAOpsPass,

      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"];
}

  %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. "