// CHECK-LABEL: private @composite_matmul_with_bias_and_relu6_fn_1
// CHECK-NEXT: %[[matmul:.*]] = "tf.MatMul"(%arg0, %arg1)
// CHECK-SAME: attr_map = "0:transpose_a,1:transpose_b"
// CHECK-NEXT: tf.BiasAdd
// CHECK-NEXT: tf.Relu6
// CHECK-NEXT: return

// CHECK-LABEL: private @composite_matmul_with_bias_and_relu_fn_1
// CHECK-NEXT: tf.MatMul"(%arg0, %arg1)
// CHECK-SAME: attr_map = "0:transpose_a,1:transpose_b"

// CHECK:  "tf.VarHandleOp"
// CHECK-NOT:  [[VARIABLE:%.*]] = "tf.ReadVariableOp"
// CHECK-NEXT:  [[KEY:%.*]], [[FUTURE:%.*]] = "tf.IfrtLoadVariable"
// CHECK-SAME:    used_by_host = true
// CHECK-NEXT:  [[MATRES:%.*]] = "tf.MatMul"(%arg0, [[FUTURE]])
// CHECK-NEXT:   [[RES:%.*]] = "tf.IfrtCall"(%arg0, [[KEY]]) <{program_id = 6515870160938153680 : i64, variable_arg_indices = [1 : i32]}>
// CHECK-NEXT:    return [[RES]], [[MATRES]] : tensor<1x1xf32>, tensor<1x1xf32>

    (TF_SubOp $beta, (TF_MulOp $m, $mul)))>;

class TFi32<int v> : ConstantAttr<I32ElementsAttr, !cast<string>(v)>;

// Matmul without transpose on b to matmul with explicit transpose op and
// transposed b.
def ConvertMatmulWithoutTransposeToWithTranspose :
      Pat<(TF_MatMulOp $a, $b, ConstBoolAttrFalse:$at, ConstBoolAttrFalse, $grad_a, $grad_b),

      ]
  )
  def matmul(self, matmul_input: core.Tensor) -> Mapping[str, core.Tensor]:
    """Performs a matrix multiplication.

    Args:
      matmul_input: Input tensor to matmul with the filter.

    Returns:
      A map of: output key -> output result.
    """
    out = math_ops.matmul(matmul_input, self.matmul_filters)

    return {'output': out}

    ).astype(np.float32)

    class TwoMatmulModel(module.Module):
      """A model with two matmul ops."""

      @def_function.function
      def matmul(self, input_tensor: core.Tensor) -> Mapping[str, core.Tensor]:
        """Performs a matrix multiplication.

        Args:
          input_tensor: Input tensor to matmul with the filter.

        Returns:
          A 'output' -> output tensor mapping
        """

    %arg1: tensor<1x3xf32> {tf_saved_model.index_path = [0]})
      -> (tensor<3x3xf32> {tf_saved_model.index_path = []}) {
  // CHECK: {{%.*}} = corert.get_op_handler %arg0 "/device:GPU:0"
  %2 = "tf.MatMul"(%arg0, %arg1) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"], device = "/device:GPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>

  // CHECK: {{%.*}} = tfrt_fallback_async.executeop {{.*}} device("/device:CPU:0") "tf.MatMul"
  // CHECK-SAME: {T = f32, transpose_a = false, transpose_b = false}

  // CHECK-NEXT: [[ch1:%.*]], [[var:%.*]] = tfrt_fallback_async.executeop.seq([[in_chain]]) {{.*}} "tf.ReadVariableOp"([[arg1]])
  // CHECK-NEXT: [[r0:%.*]] = tfrt_fallback_async.executeop {{.*}} "tf.MatMul"([[arg0]], [[var]])
  %2 = "tf.MatMul"(%arg0, %1) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"], device = "/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>

  // CHECK: %[[TRANSPOSE:.*]] = "tf.Transpose"(%[[DEQUANT]], %[[CST]]) : (tensor<3x4xf32>, tensor<?xi32>) -> tensor<*xf32>
  // CHECK: %[[MATMUL:.*]] = "tf.MatMul"(%arg0, %[[TRANSPOSE]]) <{grad_a = false, grad_b = false, transpose_a = false, transpose_b = true}> : (tensor<2x3xf32>, tensor<*xf32>) -> tensor<2x4xf32>
  // CHECK: return %[[MATMUL]] : tensor<2x4xf32>
}

  auto matmul = rewriter.create<TFL::BatchMatMulOp>(
      loc, RankedTensorType::get(matmul_shape, result_type.getElementType()),
      lhs_flattend, rhs_flattend, /*adj_x*/ false_attr, /*adj_y*/ false_attr,
      /*asym_quant_input*/ false_attr);
  if (result_type.hasStaticShape()) {
    auto reshaped =
        rewriter.create<mhlo::ReshapeOp>(loc, result_type, matmul.getResult());
    return reshaped.getResult();