%0 = "tf.BatchMatMulV2"(%arg0, %arg3) {adj_x = false, adj_y = false} : (tensor<?x8x8xf32>, tensor<8x40xf32>) -> tensor<?x8x40xf32>
  %1 = "tf.Add"(%0, %arg5) : (tensor<?x8x40xf32>, tensor<40xf32>) -> tensor<?x8x40xf32>
  %2 = "tf.BatchMatMulV2"(%1, %arg4) {adj_x = false, adj_y = true} : (tensor<?x8x40xf32>, tensor<10x40xf32>) -> tensor<?x8x10xf32>

    key: "shape"
    value {
      shape {
        dim {
          size: 3
        }
        dim {
          size: 7
        }
      }
    }
  }
}
node {
  name: "MatMul"
  op: "BatchMatMulV2"
  input: "Placeholder"
  input: "Placeholder_1"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "adj_x"
    value {
      b: false
    }
  }

    key: "shape"
    value {
      shape {
        dim {
          size: 3
        }
        dim {
          size: 7
        }
      }
    }
  }
}
node {
  name: "MatMul"
  op: "BatchMatMulV2"
  input: "Placeholder"
  input: "Placeholder_1"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "adj_x"
    value {
      b: false
    }
  }

}

INSTANTIATE_TEST_SUITE_P(
    BatchMatMulTest, BatchMatMulTest,
    ::testing::ValuesIn<MatMulTestCase>({
        {"BatchMatMul"},
        {"BatchMatMulV2"},
        {"BatchMatMulV3"},
    }),
    [](const ::testing::TestParamInfo<BatchMatMulTest::ParamType>& info) {
      return info.param.mat_mul_method;
    });

TEST(LegalizeTFTest, DumpsProducedHLO) {

  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.

// 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);

}

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

  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(