module attributes {tf.versions = {producer = 179 : i32}} {
  func.func @main(%arg0: tensor<10x19xf32>, %arg1: tensor<19x10xf32> {mhlo.is_same_data_across_replicas = true}) -> tensor<10x19xf32> {
    %0 = "tf.Shape"(%arg0) : (tensor<10x19xf32>) -> tensor<2xi64>
    %1 = "tf.Reshape"(%arg1, %0) : (tensor<19x10xf32>, tensor<2xi64>) -> tensor<10x19xf32>
    func.return %1 : tensor<10x19xf32>
  }
}

module attributes {tf.versions = {producer = 179 : i32}} {

  // CHECK-LABEL: func @main
  func.func @main(%arg0: tensor<10x19xf32>, %arg1: tensor<19x10xf32> {mhlo.is_same_data_across_replicas = true}) -> tensor<?xi64> {
    %0 = "tf.Shape"(%arg0) : (tensor<10x19xf32>) -> tensor<2xi64>
    %1 = "tf.Reshape"(%arg1, %0) : (tensor<19x10xf32>, tensor<2xi64>) -> tensor<?x?xf32>

    // CHECK: %[[RESULT:.*]] = mhlo.constant dense<[10, 19]>

}

func.func @matmul_batchv3_unknown_dim(%arg0: tensor<?x10x15xf32>, %arg1: tensor<15x17xf32>) -> tensor<?x10x17xf32> {
  %0 = "tf.BatchMatMulV3"(%arg0, %arg1) {Ta = "tfdtype$DT_FLOAT", Tb = "tfdtype$DT_FLOAT", device = "/device:CPU:0", name = "MatMul", adj_x = false, adj_y = false} :
(tensor<?x10x15xf32>, tensor<15x17xf32>) -> tensor<?x10x17xf32>

//...

      /*error_tolerance=*/1.0);
}

TEST_F(ConvertTfQuantToMhloIntTest, UniformQuantizeAddToValidGraph) {
  constexpr absl::string_view kProgram = R"mlir(
func.func @main(%lhs: tensor<10x10xi32>, %rhs: tensor<10x10xi32>) -> tensor<10x10xi32> {
  %lhs_scale = "tf.Const"() { value = dense<0.518> : tensor<f32> } : () -> tensor<f32>
  %lhs_zp = "tf.Const"() { value = dense<42> : tensor<i32> } : () -> tensor<i32>

  %1 = "tf.MatMul"(%arg0, %arg0) {device = "", transpose_a = false, transpose_b = false} : (tensor<10x10xi32>, tensor<10x10xi32>) -> tensor<10x10xi32>
  %2 = "tf.PreventGradient"(%0) : (tensor<10x10xi32>) -> tensor<10x10xi32>
  %3 = "tf.PreventGradient"(%1) : (tensor<10x10xi32>) -> tensor<10x10xi32>
  %4 = "tf.AddV2"(%2, %3) {device = ""} : (tensor<10x10xi32>, tensor<10x10xi32>) -> tensor<10x10xi32>

// CHECK-SAME: (%[[ARG0:.*]]: tensor<10x8x9xf32>, %[[ARG1:.*]]: tensor<5xi32>, %[[ARG2:.*]]: tensor<5x8x9xf32>) -> tensor<10x8x9xf32>
func.func @main(%arg0: tensor<10x8x9xf32>, %arg1: tensor<5xi32>, %arg2: tensor<5x8x9xf32>) -> tensor<10x8x9xf32> {
  %elem_shape = "tf.Const"() {value = dense<[8, 9]> : tensor<2xi32>} : () -> tensor<2xi32>
  // CHECK: %[[BUFFER:.*]] = "tf.Identity"(%[[ARG0]]) : (tensor<10x8x9xf32>) -> tensor<10x8x9xf32>

  %0 = arith.constant dense<-1> : tensor<1xi32>
  %1 = "tfl.reduce_max"(%arg0, %0) {keep_dims = true} : (tensor<10x10xf32>, tensor<1xi32>) -> tensor<10x1xf32>
  %2 = tfl.sub(%arg0, %1) {fused_activation_function = "NONE"} : (tensor<10x10xf32>, tensor<10x1xf32>) -> tensor<10x10xf32>
  %3 = "tfl.exp"(%2) : (tensor<10x10xf32>) -> tensor<10x10xf32>

func.func @testSoftmax(%arg0: tensor<10x10xf32>) -> tensor<10x10xf32> {
  // CHECK: _arithmetic_count = 6400 : i64
  %0 = "tfl.softmax"(%arg0) {beta = 1.000000e+00 : f32} : (tensor<10x10xf32>) -> tensor<10x10xf32>
  func.return %0 : tensor<10x10xf32>
}

func.func @testTanh(%arg0: tensor<10x10xf32>) -> tensor<10x10xf32> {
  // CHECK: _arithmetic_count = 6400 : i64
  %0 = "tfl.tanh"(%arg0) : (tensor<10x10xf32>) -> tensor<10x10xf32>

  func.return %0 : tensor<10x10xf32>