func.return %0 : tensor<4xf32>
}

// CHECK:       func @squaredDifference(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
// CHECK:         %0 = "tf.Sub"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
// CHECK:         %1 = "tf.Mul"(%0, %0) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
// CHECK:         return %1 : tensor<4xf32>
// CHECK:       }

// -----

// CHECK: return %[[RESULT0]]
}

// CHECK-LABEL: testMatmulAsymAttributeTrue
func.func @testMatmulAsymAttributeTrue(%arg0: tensor<4x10x15xf32>, %arg1: tensor<4x15x17xf32>) -> tensor<4x10x17xf32> {
  %0 = "tfl.batch_matmul"(%arg0, %arg1) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = true} : (tensor<4x10x15xf32>, tensor<4x15x17xf32>) -> tensor<4x10x17xf32>

}

// CHECK-LABEL: testAddN
func.func @testAddN(tensor<? x f32>, tensor<? x f32>, tensor<? x f32>) -> tensor<? x f32> {
^bb0(%arg0: tensor<? x f32>, %arg1: tensor<? x f32>, %arg2: tensor<? x f32>):
  // CHECK: "tfl.add_n"(%arg0, %arg1, %arg2)
  %0 = "tfl.add_n"(%arg0, %arg1, %arg2): (tensor<? x f32>, tensor<? x f32>, tensor<? x f32>) -> tensor<? x f32>
  func.return %0 : tensor<? x f32>
}

// -----

    "tfl.yield"(%6, %5, %32, %17, %20, %0) : (tensor<i32>, tensor<i32>, tensor<*xf32>, tensor<4x2xf32>, tensor<4x2xf32>, tensor<4x4x3xf32>) -> ()
  }) {is_stateless = true} : (tensor<i32>, tensor<i32>, tensor<4x4x2xf32>, tensor<4x2xf32>, tensor<4x2xf32>, tensor<4x4x3xf32>) -> (tensor<i32>, tensor<i32>, tensor<*xf32>, tensor<4x2xf32>, tensor<4x2xf32>, tensor<*xf32>)

    R"mlir(
  func.func @main(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
    %if = "stablehlo.if"(%arg0) ({
      %0 = stablehlo.add %arg1, %arg1 : tensor<f32>
      stablehlo.return %0 : tensor<f32>
    }, {
      %1 = stablehlo.add %arg2, %arg2 : tensor<f32>
      stablehlo.return %1 : tensor<f32>
    }) : (tensor<i1>) -> (tensor<f32>)

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

// CHECK-LABEL: func @string_type
func.func @string_type(%arg: tensor<1x2x!tf_type.string>) -> tensor<?x6x!tf_type.string> {
  %multiples = "tf.Const"() { device = "/device:CPU:0", value = dense<[7,3]> : tensor<2xi32> } : () -> tensor<2xi32>
  // CHECK: T = !corert.string

  %2 = "tf.Mul"(%arg1, %arg3) : (tensor<?x1xf32>, tensor<f32>) -> tensor<?x1xf32>
  // Add of a scalar const and a tensor.
  %3 = "tf.AddV2"(%arg2, %1) : (tensor<f32>, tensor<?x1xf32>) -> tensor<?x1xf32>
  %4 = "tf.AddV2"(%arg3, %2) : (tensor<f32>, tensor<?x1xf32>) -> tensor<?x1xf32>

// CHECK:         }
func.func @tfl_wrapped_jax_random_normal(%arg0: tensor<2xui32>) -> tuple<tensor<3x4xf32>> {
  // This is a fake jax random normal body.
  %0 = stablehlo.constant dense<0.0> : tensor<12xf32>
  %1 = "stablehlo.reshape"(%0) : (tensor<12xf32>) -> tensor<3x4xf32>
  %2 = "stablehlo.tuple"(%1) : (tensor<3x4xf32>) -> tuple<tensor<3x4xf32>>
  func.return %2 : tuple<tensor<3x4xf32>>
}

  %9 = "tfl.cos"(%2) : (tensor<f32>) -> tensor<f32>
  %10 = "tfl.log"(%3) : (tensor<f32>) -> tensor<f32>
  %11 = "tfl.sqrt"(%4) : (tensor<f32>) -> tensor<f32>
  %12 = "tfl.rsqrt"(%5) : (tensor<f32>) -> tensor<f32>
  %13 = "tfl.square"(%6) : (tensor<f32>) -> tensor<f32>

  func.return %7, %8, %9, %10, %11, %12, %13 : tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>
}

  func.return %2 : tensor<1x384x32xf32>
}

// PARTIAL-LABEL: @simple_test
func.func @simple_test(%arg0: tensor<4x384x32xf32>, %arg1: tensor<3xi32>, %arg2: tensor<3xi32>) -> tensor<1x384x32xf32> attributes {tac.interface_name = "func1"} {
  %0 = "tfl.slice"(%arg0, %arg1, %arg2) : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>