}

// -----
// ==== V2 tests ====

func.func @batchMatMulV2TwoDim(%arg0: tensor<2x3x4x5xf32>, %arg1: tensor<2x3x5x6xf32>) -> tensor<2x3x4x6xf32> {
  %0 = "tf.BatchMatMulV2"(%arg0, %arg1) : (tensor<2x3x4x5xf32>, tensor<2x3x5x6xf32>) -> tensor<2x3x4x6xf32>
  func.return %0 : tensor<2x3x4x6xf32>

  // CHECK-LABEL: batchMatMulV2TwoDim

    %7 = "tf.Identity"(%cst_2) {device = ""} : (tensor<20x30x1024x3xi8>) -> tensor<20x30x1024x3xi8>
    %8 = "tf.Cast"(%7) {Truncate = false, device = ""} : (tensor<20x30x1024x3xi8>) -> tensor<20x30x1024x3xi32>
    %9 = "tf.BatchMatMulV2"(%6, %8) {adj_x = false, adj_y = false, device = ""} : (tensor<20x30x64x1024xi32>, tensor<20x30x1024x3xi32>) -> tensor<20x30x64x3xi32>

  (*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,

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