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

  auto matmul = rewriter.create<TF::BatchMatMulV3Op>(
      loc, RankedTensorType::get(matmul_shape, result_type.getElementType()),
      lhs_flattend, rhs_flattend);

  if (result_type.hasStaticShape()) {
    auto reshaped =
        rewriter.create<mhlo::ReshapeOp>(loc, result_type, matmul.getResult());
    return reshaped.getResult();
  }

  return %1#0, %1#1 : tensor<1x2xf32>, tensor<1x2xf32>
}
func.func @_func(%arg0: tensor<2x4xf32>, %arg1: tensor<4x2xf32>) -> tensor<2x2xf32> {
  %0 = "tf.MatMul"(%arg0, %arg1) {_XlaSharding = "\08\03\1A\02\02\01\22\02\00\01"} : (tensor<2x4xf32>, tensor<4x2xf32>) -> tensor<2x2xf32>
  return %0 : tensor<2x2xf32>
}

// -----

foreach src = [TF_PreventGradientOp, TF_CheckNumericsOp] in
  def : Pat<(src $op, $msg), (replaceWithValue $op)>;

//===----------------------------------------------------------------------===//
// MatMul op patterns.
//===----------------------------------------------------------------------===//

def GetPrecisionConfig: NativeCodeCall<
  "GetPrecisionConfig(&$_builder)">;

  return op;
}

TFE_Op* MatMulOp(TFE_Context* ctx, TFE_TensorHandle* a, TFE_TensorHandle* b) {
  TF_Status* status = TF_NewStatus();

  TFE_Op* op = TFE_NewOp(ctx, "MatMul", status);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TFE_OpAddInput(op, a, status);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TFE_OpAddInput(op, b, status);

// Guarantee that all FuncOp's have a single use.
std::unique_ptr<OperationPass<ModuleOp>> CreateGuaranteeAllFuncsOneUsePass();

// Optional pass which will unroll BatchMatMul and use only MatMul
std::unique_ptr<OperationPass<func::FuncOp>> CreateUnrollBatchMatMulPassPass();

// Optional pass which will map TF BatchMatMul to TF Einsum
std::unique_ptr<OperationPass<func::FuncOp>> CreateBatchMatMulToEinsumPass();

//   %1 = mhlo.reshape %param : (1xCxZ) -> CxZ
//   mhlo.dot_general %input, %1 {batch_dims = []}
// To:
//   mhlo.dot_general %input, %param {batch_dims = [0]}
//
// This usage will mostly come from tf-unroll-batch-matmul, so it's fine to only
// handle the case where batching dim is the leftmost dim.
LogicalResult ConvertReshapeDotRhsToBatchedDot(mhlo::DotGeneralOp dot,
                                               PatternRewriter &rewriter) {

    //
    // 2. Old fashion Tensor with raw device memory pointer. This case occurs
    // when the producer is a non-XLA TF GPU kernel or function (e.g.
    // tf.matmul).
    //
    // 3. AsyncValueTensor, containing a PjRtBuffer. This is the legacy mode
    // and certain device type (e.g. TPU) still uses this path.
    AsyncValueTensor* av_tensor = AsyncValueTensor::FromTensor(tensor);

    quantization_axis = -1 : i64, quantization_min_val = -128 : i64,
    quantization_max_val = 127 : i64
  } : (tensor<9x10x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<9x10xf32>
  %0 = "tf.MatMul"(%input, %filter_new) {
  } : (tensor<8x9xf32>, tensor<9x10xf32>) -> tensor<8x10xf32>
  return %0 : tensor<8x10xf32>
})mlir";
  constexpr absl::string_view kProgram = R"mlir(

func.func @FuseMulWithFullyConnectedWithBias(%arg: tensor<2x512xf32>) -> tensor<2x1024xf32> {
  %cst_mul = arith.constant dense<2.0> : tensor<512xf32>
  %cst_weights = arith.constant dense<3.0> : tensor<1024x512xf32>
  %cst_bias = arith.constant dense<5.0> : tensor<1024xf32>

  %0 = "tfl.mul"(%arg, %cst_mul) {fused_activation_function = "NONE"} : (tensor<2x512xf32>, tensor<512xf32>) -> tensor<2x512xf32>