// CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_4]], %[[RHS_4]]) <{grad_a = false, grad_b = false, transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32>

def ConvertTFMatMulToXLADotV2Op : Pat<
  (TF_MatMulOp:$matmul
    (TF_SubOp (TF_CastOp $input, $truncate), $input_zp),
    (TF_CastOp (TF_IdentityOp $weight), $truncate1),
    $transpose_a, $transpose_b, $grad_a, $grad_b),
  (CreateXlaDotV2OpFromTfMatMulOp
    $input, $weight, $input_zp,
    /*weight_zp=*/(CreateScalarIntegerConst<"int32_t", "0">), $matmul,
    $transpose_a, $transpose_b),
  [(IsInt8ElementType $input),

      (NamedAttr<"dilations"> $dilations))),
  [(IsNotInLiftedFunc $res)], [], (addBenefit 1)>;

def LiftMatMul : Pat<
  (TF_MatMulOp:$res $a, $b, $transpose_a, $transpose_b, $grad_a, $grad_b),
  (LiftAsTFPartitionedCall<"composite_matmul_fn">
    (ArgumentList $a, $b),
    (ResultList $res),
    (NamedAttributeList
      (NamedAttr<"transpose_a"> $transpose_a),

  TF_RETURN_IF_ERROR(op_ptr->SetAttrBool("transpose_a", transpose_a));
  TF_RETURN_IF_ERROR(op_ptr->SetAttrBool("transpose_b", transpose_b));
  TF_RETURN_IF_ERROR(op_ptr->SetAttrBool("grad_a", grad_a));
  TF_RETURN_IF_ERROR(op_ptr->SetAttrBool("grad_b", grad_b));
  int num_retvals = 1;
  return op_ptr->Execute(absl::MakeSpan(product, 1), &num_retvals);
}

// Op: IdentityN()
// Summary:
//
// Description:

      Pat<(TF_MatMulOp $a, $b, ConstBoolAttrFalse:$at, ConstBoolAttrFalse, $grad_a, $grad_b),
          (TF_MatMulOp $a, (TF_TransposeOp $b, (TF_SubOp (TF_RangeOp
             /*start=*/(TF_RankOp $b),
             /*limit=*/(TF_ConstOp TFi32<0>),
             /*delta=*/(TF_ConstOp TFi32<-1>)), (TF_ConstOp TFi32<1>))),
           $at, ConstBoolAttrTrue, $grad_a, $grad_b)>;

// Matmul with transpose on a to matmul with explicit transpose op and a not

      });
      return false;
    }
    // FusedMatMul kernel does not support grad_a/grad_b attrs
    if ((matmul->hasAttr("grad_a") &&
         mlir::cast<BoolAttr>(matmul->getAttr("grad_a")).getValue()) ||
        (matmul->hasAttr("grad_b") &&
         mlir::cast<BoolAttr>(matmul->getAttr("grad_b")).getValue())) {
      (void)rewriter.notifyMatchFailure(matmul, [&](Diagnostic &diag) {

//
Status MatMul(AbstractContext* ctx, AbstractTensorHandle* const a, AbstractTensorHandle* const b, AbstractTensorHandle** product, bool transpose_a = false, bool transpose_b = false, bool grad_a = false, bool grad_b = false, const char* name = nullptr, const char* raw_device_name = nullptr);

//

  updates_grad_reshaped = tf.reduce_sum(
      grad, axis=reduction_axes, keepdims=True)
  bias_grad = tf.reshape(updates_grad_reshaped, input_value_shape)

  a = math_ops.conj(op.inputs[0])
  b = math_ops.conj(op.inputs[1])
  grad_a = gen_math_ops.mat_mul(grad, b)
  grad_b = gen_math_ops.mat_mul(grad, a, transpose_a=True)
  return [grad_a, grad_b, bias_grad]


@Composite(
    'NewMaxPool',

  }

  func.func private @composite_matmul_with_bias_fn_1(%arg0: tensor<1x4xf32>, %arg1: tensor<4x3xf32>, %arg2: tensor<3xf32>) -> tensor<1x3xf32> attributes {tf_quant.composite_function} {
    %0 = "tf.MatMul"(%arg0, %arg1) <{grad_a = false, grad_b = false, transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_b", device = ""} : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>

      (NamedAttr<"dilations"> $dilations))),
  [(IsNotInLiftedFunc $res), (IsConstTensor $filter)], [], (addBenefit 1)>;

def LiftMatMul : Pat<
  (TF_MatMulOp:$res $a, $b, $transpose_a, $transpose_b, $grad_a, $grad_b),
  (LiftAsTFPartitionedCall<"composite_matmul_fn">
    (ArgumentList $a, $b),
    (ResultList $res),
    (NamedAttributeList
      (NamedAttr<"transpose_a"> $transpose_a),