func.func private @composite_matmul_fn(%arg0: tensor<1x2x2x3xf32>, %arg1: tensor<2x1024xf32>) -> tensor<*xf32> attributes {tf_quant.composite_function} {
    %0 = "tf.MatMul"(%arg0, %arg1) {attr_map = "0:transpose_a,1:transpose_a", device = "", transpose_a = false, transpose_b = false} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32>
    return %0 : tensor<*xf32>
  }

  // CHECK: device = "xpu"
  %1 = "tf.MatMul"(%arg0, %0) {T = f32, _output_shapes = ["tfshape$dim { size: 3 } dim { size: 3 }"], device = "", transpose_a = false, transpose_b = false} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>
  // CHECK: device = "cpu"

                                              /*a=*/input_lhs,
                                              /*b=*/input_rhs,
                                              /*transpose_a=*/op.getAdjX(),
                                              /*transpose_b=*/op.getAdjY());
    return success();
  }

  // Input dimensions must be defined. MatMulBCast does not support partial
  // shapes.
  for (auto dim : lhs_shape) {

    (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),

  %0 = "tf.MatMul"(%arg0, %arg1) {
    transpose_a = false, transpose_b = false
  } : (tensor<1x10xf32>, tensor<10x10xf32>) -> tensor<*xf32>
  %1 = "tf.BiasAdd"(%0, %cst) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<10xf32>) -> tensor<*xf32>
  %2 = "tf.Relu6"(%1) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>

  %3 = "tf.MatMul"(%arg0, %arg1) {
    transpose_a = true, transpose_b = false

//     .Output("product: T")
//     .Attr("transpose_a: bool = false")
//     .Attr("transpose_b: bool = false")
//     .Attr("T: {bfloat16, half, float, double, int32, int64, complex64, complex128}")
// T is a derived attribute.
// transpose_a and transpose_b is materialized attributes.
tfr.func @tf__mat_mul_(!tfr.tensor<T>, !tfr.tensor<T>,
                      i1 {tfr.name="transpose_a", tfr.default=false},

      auto dz = Const(scope, {{1.0, 1.0}, {1.0, 1.0}});
      auto dx = MatMul(scope, dz, y, MatMul::TransposeB(true));
      auto dy = MatMul(scope, x, dz, MatMul::TransposeA(true));

      auto du = MatMul(scope, dx, v, MatMul::TransposeB(true));
      auto dv = MatMul(scope, u, dx, MatMul::TransposeA(true));
    } else {
      // Call AddSymbolicGradients.
      auto dz = Const(scope, {{1.0, 1.0}, {1.0, 1.0}});

  func.func private @composite_matmul_fn(%arg0: tensor<1x2x2x3xf32>, %arg1: tensor<2x1024xf32>) -> tensor<*xf32> attributes {tf_quant.composite_function} {
    %0 = "tf.MatMul"(%arg0, %arg1) {attr_map = "0:transpose_a,1:transpose_a", device = "", transpose_a = false, transpose_b = false} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32>
    return %0 : tensor<*xf32>
  }

// CHECK-LABEL: func @matmul

  func.func private @composite_matmul_fn_2(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> attributes {tf_quant.composite_function} {
    %0 = "tf.MatMul"(%arg0, %arg1) {attr_map = "0:transpose_a,1:transpose_b", device = "", transpose_a = false, transpose_b = false} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
    return %0 : tensor<2x2xf32>
  }

Aditya Manthramurthy <******@****.***> 1716591923 -0700