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',
    inputs=['input_: T'],
    attrs=[

   (IsInt8ElementType $weight),
   (IsConstTensor $weight),
   (IsInt32ElementType $matmul),
   (HasStaticShapeConstraint $weight)],
  [], (addBenefit 10)>;

// Convert Matmul with hybrid inputs (f32 activation/int8 weight) to XlaDotV2
def ConvertTFMatMulToXLADotV2OpWeightOnly : Pat<
  (TF_MatMulOp:$matmul
    $input,
    (TF_MulOp (TF_CastOp (TF_IdentityOp $weight), $truncate1), $scale),

  input: "sequential/quant_dense/MatMul/kquant/FakeQuantWithMinMaxVars/ReadVariableOp"
  input: "sequential/quant_dense/MatMul/kquant/FakeQuantWithMinMaxVars/ReadVariableOp_1"
  attr {
    key: "narrow_range"
    value {
      b: false
    }
  }
  attr {
    key: "num_bits"
    value {
      i: 8
    }
  }
}
node {
  name: "sequential/quant_dense/MatMul/kquant/IdentityN"

# CHECK-DAG: %[[MUL1:.*]] = "tf.MatMul"(%[[ARG0]], %[[ARG1]])
# CHECK-DAG: %[[MUL2:.*]] = "tf.MatMul"(%[[ARG1]], %[[ARG0]])
# CHECK:  %[[IDENTITY:.*]]:2 = "tf.IdentityN"(%[[MUL1]], %[[MUL2]])
# CHECK: return %[[IDENTITY]]#1, %[[IDENTITY]]#0


def Test():

  x = tf.constant(1.0, shape=(5, 3))
  y = tf.constant(1.0, shape=(3, 5))

  s = tf.matmul(x, y)
  t = tf.matmul(y, x)
  [t, s] = array_ops.identity_n([t, s])

  op: "Identity"
  input: "Placeholder_1"
  attr {
    key: "T"
    value {
      type: DT_BOOL
    }
  }
}
node {
  name: "cond/MatMul"
  op: "MatMul"
  input: "cond/MatMul/Switch:1"
  input: "cond/MatMul/Switch_1:1"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "transpose_a"
    value {
      b: false
    }

    %outputs_7, %control_8 = tf_executor.island wraps "tf.Const"() {value = dense<"MatMul/b_0"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
    %outputs_9, %control_10 = tf_executor.island wraps "tf.VarHandleOp"() {container = "", shared_name = "MatMul/b_0"} : () -> tensor<!tf_type.resource<tensor<20x4096xf32>>>

# RUN: tf_tfl_translate -unfold_batchmatmul=false -tf-input-arrays=Placeholder,Placeholder_1 -tf-input-shapes=2,5,3:3,7 -tf-input-data-types=DT_FLOAT,DT_FLOAT -tf-output-arrays=MatMul -output-mlir %s -o - 2>&1 | FileCheck %s

node {
  name: "Placeholder"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "shape"
    value {
      shape {
        dim {
          size: 2

 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

  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

  TFE_Op* matmul = MatMulOp(ctx, h0_task1, h1_task1);
  TFE_OpSetDevice(matmul, remote_device_name, status);
  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

  TFE_TensorHandle* retvals[1];
  int num_retvals = 1;
  TFE_Execute(matmul, &retvals[0], &num_retvals, status);
  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

  }
  return ConstantFoldOpIfPossible(value.getDefiningOp()).front();
}

// Matches convolution op with "NHWC" data format or matmul op with false adj_y.
// The list of supported ops in this function is:
// - Conv2DOp
// - Conv3DOp
// - DepthwiseConv2dNativeOp
// - MatMulOp
// - BatchMatMulV2Op
LogicalResult MatchSupportedAffineOp(Operation* op, Value& binding_output,