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=[

  s = s.replace('2D', '2d').replace('3D', '3d')
  snake_case = ''.join(['_' + i.lower() if i.isupper() else i for i in s
                       ]).lstrip('_')
  return snake_case.replace('mat_mul', 'matmul').replace('bias_add', 'bias')


def _substitute_impl_function_name_template(module: str) -> str:
  """Generates the op-specific implementation function name."""

Christian Sigg <******@****.***> 1714640622 -0700

      "/job:localhost/replica:0/task:0/device:CPU:0");
  Output a = ops::Placeholder(root.WithOpName("A"), DT_FLOAT);
  for (int i = 0; i < 5; i++) {
    a = ops::MatMul(root.WithOpName(absl::StrCat("matmul_", i)), a, a);
    a = ops::Add(root.WithOpName(absl::StrCat("add_", i)), a, a);
  }

  GraphDef graph_def;
  root.graph()->ToGraphDef(&graph_def);
  return graph_def;
}

func.func @fold_identity_n_test(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> (tensor<2x2xf32>, tensor<2x2xf32>) {
  // CHECK: tf.MatMul
  %outputs = "tf.MatMul"(%arg0, %arg1) {device = "TPU", transpose_a = false, transpose_b = false} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
  %outputs_0 = "tf.MatMul"(%arg0, %arg1) {device = "TPU", transpose_a = false, transpose_b = false} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>

    TFE_OpAddInput(matmul, h0_task0, status);
    ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
    TFE_OpAddInput(matmul, has_packed_input ? packed_handle : h1_task2, status);
    ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
  } else {
    // Handles are on task0 (local), and task2, but op is on task1.
    matmul = MatMulOp(ctx, h0_task0, h1_task2);
  }
  if (remote) {

# 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
    }

# 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