Transpose 2
 Unpack 1
cluster 2 size 44
 Cast 2
 ConcatV2 2
 Const 18
 ExpandDims 1
 GatherV2 2
 Less 1
 MatMul 1
 Mul 1
 Pack 1
 Prod 2
 Range 1
 Reshape 3
 Shape 8
 StridedSlice 1
cluster 3 size 10
 AddN 1
 Const 1
 MatMul 2
 Mul 1
 Reshape 3
 Sum 1
 Transpose 1
cluster 4 size 11
 ConcatOffset 1
 Const 4
 ReverseSequence 1
 Slice 2

    std::string name_grad_B = "MatMul_Grad_B";
    if (!t_a && !t_b) {
      TF_RETURN_IF_ERROR(MatMul(ctx, upstream_grad, B.get(), &matmul_A_output,
                                /*transpose_a = */ false,
                                /*transpose_b = */ true, name_grad_A.c_str()));

      TF_RETURN_IF_ERROR(MatMul(ctx, A.get(), upstream_grad, &matmul_B_output,
                                /*transpose_a = */ true,

// CHECK: %[[IDENTITY:.*]] = "tf.Identity"(%cst) : (tensor<2x1024xi8>) -> tensor<2x1024xi8>
// CHECK: %[[MATMUL:.*]] = "tf.XlaDotV2"(%arg0, %[[IDENTITY]]) <{dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""}> {device = ""} : (tensor<1x2x2x2xbf16>, tensor<2x1024xi8>) -> tensor<1x2x2x1024xbf16>

// RUN: tf-quant-opt %s -split-input-file -quant-lift-quantizable-spots-as-functions -quant-prepare-quantize-drq -quant-quantize='weight-quantization=true' -verify-each=false | FileCheck %s

// -----

module {
  func.func @matmul(%arg0: tensor<1x2x2x3xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<2x1024xf32>} : () -> tensor<2x1024xf32>

  // CHECK: tfrt_fallback_async.executeop key(2) cost({{.*}}) device("/device:CPU:0") "tf.MatMul"
  %0 = "tf.ReadVariableOp"(%arg1) {device = "/device:CPU:0", dtype = f32} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
  %1 = "tf.MatMul"(%arg0, %0) {T = f32, device = "/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>

  %cst = "tf.Const"() {value = dense<0.000000e+00> : tensor<512x512xf32>} : () -> tensor<512x512xf32>
  %out_1 = "tf.MatMul"(%arg0, %cst) {
    device = "", transpose_a = false, transpose_b = false
  } : (tensor<1x12x12x512xf32>, tensor<512x512xf32>) -> tensor<*xf32>
  %out_2 = "tf.MatMul"(%arg0, %arg0) {
    device = "", transpose_a = false, transpose_b = true
  } : (tensor<1x12x12x512xf32>, tensor<1x12x12x512xf32>) -> tensor<*xf32>

module {
  // TODO(b/260020937): Support transpose_a, transpose_b for matmul.
  func.func @matmul(%arg0: tensor<2x12xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<12x2xf32>} : () -> tensor<12x2xf32>

  %out_matrix = "tf.MatMul"(%in_matrix, %matrix)  : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
  %in_matrix1 = "tf.TensorArrayReadV3"(%handle_2, %loop_count, %flow_in_2) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
  %out_matrix1 = "tf.MatMul"(%out_matrix, %matrix_2)  : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>

category names correspond to generated source file names, and should be
consistent with the original source files registering each operator. For example
since `REGISTER_OP("MatMul")` appears in ***core/math_ops.cc***, the "MatMul"
operator in the script should be in the "math" category, and it will be
generated in the output file `c/experimental/ops/math_ops.cc`.

# CHECK-NEXT: [[R1:%.*]] = "tf.ReadVariableOp"([[ARG1]]) {{{.*}}} : (tensor<!tf_type.resource<tensor<3x5xf32>>>) -> tensor<3x5xf32>
# CHECK-NEXT: [[R2:%.*]] = "tf.MatMul"([[R0]], [[R1]]) <{{{.*}}}> {{{.*}}} : (tensor<5x3xf32>, tensor<3x5xf32>) -> tensor<5x5xf32>


def Test():

  x = tf.compat.v1.get_variable(
      name='x',
      shape=(5, 3),