%2 = "tf.Cast"(%identity) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
    %3 = "tf.Sub"(%2, %weight_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>

    %5 = "tf.BatchMatMulV2"(%1, %3) {
      attr_map = "adj_x:0,adj_y:1"
    } : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
    func.return %5 : tensor<*xi32>
  }

   // Einsum with int32 accumulation.

  for (int64_t i = 0; i < input.size(); ++i) {
    output[permutation[i]] = input[i];
  }
  return output;
}

// Computes the transpositions required to convert dnums to one supported by
// tf.BatchMatmulV2 and returns the new set of dimension numbers with them.
// Transposed LHS shape will be B0,...,Bn,L0,...,Ln,C0,...,Cn and,
// transposed RHS shape will be B0,...,Bn,C0,...,Cn,R0,...,Rn respectively.

  func.return %0 : tensor<3x4x6xf32>
}
// CHECK-LABEL: lower_einsum
// CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) <{adj_x = false, adj_y = false}> : (tensor<3x4x5xf32>, tensor<3x5x6xf32>) -> tensor<3x4x6xf32>

// -----

func.func @removing_identity_after_const(%arg0: tensor<*xf32>) -> (tensor<*xf32>) {

// operands are properly supported in declarative rewrite rule specification.

DECL_CONVERT_OP(Assert);
DECL_CONVERT_OP(ConcatV2);
DECL_CONVERT_OP(BatchMatMul);
DECL_CONVERT_OP(BatchMatMulV2);
DECL_CONVERT_OP(BatchMatMulV3);
DECL_CONVERT_OP(MatMul);
DECL_CONVERT_OP(MatrixDiagV2);
DECL_CONVERT_OP(MatrixDiagV3);
DECL_CONVERT_OP(Pack);
DECL_CONVERT_OP(Split);
DECL_CONVERT_OP(SplitV);

  (*grad_outputs)[1] =
      Reshape(scope, ReduceSum(scope, (*grad_outputs)[1], reduce.r1), sy);
  return scope.status();
}
REGISTER_GRADIENT_OP("BatchMatMulV2", BatchMatMulV2Grad);
REGISTER_GRADIENT_OP("BatchMatMulV3", BatchMatMulV2Grad);

Status CumsumGrad(const Scope& scope, const Operation& op,
                  const std::vector<Output>& grad_inputs,

}

func.func @matmul_batchv2(%arg0: tensor<2x10x15xf32>, %arg1: tensor<15x17xf32>) -> tensor<2x10x17xf32> {
  %0 = "tf.BatchMatMulV2"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", device = "/device:CPU:0", name = "MatMul", adj_x = false, adj_y = false} :
(tensor<2x10x15xf32>, tensor<15x17xf32>) -> tensor<2x10x17xf32>
  func.return %0 : tensor<2x10x17xf32>

  func.return %0: tensor<2x3xf32>
}

// CHECK-LABEL: testBatchMatMulToV2
func.func @testBatchMatMulToV2(%arg0: tensor<2x3x5xf32>, %arg1: tensor<2x5x7xf32>) -> tensor<2x3x7xf32> {
  // CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) <{adj_x = false, adj_y = false, grad_x = false, grad_y = false}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"}

      "AssignVariableOp",
      "AssignVariableXlaConcatND",
      "AvgPool",
      "AvgPool3D",
      "AvgPool3DGrad",
      "AvgPoolGrad",
      "BatchMatMul",
      "BatchMatMulV2",
      "BatchMatMulV3",
      "BatchToSpace",
      "BatchToSpaceND",
      "BesselI0e",
      "BesselI1e",
      "Betainc",
      "BiasAddV1",
      "Bincount",
      "Bucketize",

  llvm::ArrayRef<int64_t> y_batches = y_shape.drop_back(2);

  // Check compatibility of batch dimensions if both input shapes are known.
  // BatchMatMul should have exactly the same batch dimensions and
  // BatchMatMulV2 should have broadcastable batch dimensions.
  //
  // The last two dimensions are non-batch dimensions that don't need to
  // participate in batch dimension compatibility check.
  if (std::is_same<OpT, BatchMatMulOp>()) {

    if target_opset == quant_opts_pb2.UNIFORM_QUANTIZED:
      self.assertFalse(self._contains_op(output_graphdef, 'XlaDotV2'))
      self.assertTrue(self._contains_op(output_graphdef, 'BatchMatMulV2'))
    else:
      self.assertFalse(self._contains_op(output_graphdef, 'XlaDotV2'))
      self.assertTrue(self._contains_op(output_graphdef, 'Einsum'))

  @parameterized.named_parameters(