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."""

  // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) <{grad_a = false, grad_b = false, transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32>
  // 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>

  // done in parallel.
  //
  // This graph is:
  // (Const0, Const0) -> MatMul0
  // (Const1, Const1) -> MatMul1
  // (MatMul0, MatMul1) -> MatMulCombined
  //
  // Device0: [Const0, Const0, MatMul0]
  // Device1: [Const1, Const1, MatMul1, MatMulCombined]
  //
  // Cluster0: [Const0, Const0, MatMul0]
  // Cluster1: [Const1, Const1, MatMul1]
  // Cluster2: [MatMulCombined]

// TF-DAG: %[[pc_3:.*]] = "tf.PartitionedCall"(%arg0, %[[cst_1]]) <{config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn_2_0}

// IntPerLayer-DAG: "tf.DumpTensor"(%[[m0_1]]) <{enabled = true, file_name = "unquantized_tensor_data.pb", func_name = "matmul2", log_dir_path = "/tmp/dumps/composite_matmul_fn_2", node_name = "MatMul"}> : (tensor<2x2xf32>) -> ()

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

                                                /*transpose_b=*/op.getAdjY());
    matmuls.emplace_back(matmul.getProduct());
  }

  // Combine the result of each individual MatMul into a rank-3 tensor.
  Type packed_type = RankedTensorType::get(
      {bcast.output_batch_size(), rows, cols}, element_type);
  const auto axis = rewriter.getI64IntegerAttr(0);
  auto pack_op =

  TFE_TensorHandle* m = TestMatrixTensorHandle(ctx);
  TFE_Op* matmul = TFE_NewOp(ctx, "MatMul", status);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TFE_TensorHandle* retvals[1];
  int num_retvals = 1;
  for (auto s : state) {
    TFE_OpReset(matmul, "MatMul", nullptr, status);
    CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
    TFE_OpAddInput(matmul, m, status);

      "/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;
}