} : (tensor<10xf32>, tensor<10xf32>, tensor<6xi32>) -> (tensor<10xf32>, tensor<6xi32>)
    // CHECK: return %[[RESULT_CAST]], %[[RESULT]]#1 : tensor<10xf32>, tensor<6xi32>
    func.return %0#0, %0#1 : tensor<10xf32>, tensor<6xi32>
  }

  // CHECK-LABEL: func private @main_0
  // CHECK-SAME: %[[ARG_0:.*]]: tensor<10xbf16>, %[[ARG_1:.*]]: tensor<10xbf16>, %[[ARG_2:.*]]: tensor<6xi32>
  func.func private @main_0(

    %3 = "tf.Mul"(%1, %2) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
    return %3 : tensor<10xf32>
  }
}

// -----

// CHECK-LABEL: module
module attributes {tf_saved_model.semantics} {
  // CHECK: ml_program.global{{.*}}mutable{{.*}}@vars.v
  "tf_saved_model.global_tensor"() { is_mutable, sym_name = "v", type = tensor<10xf32>, value = dense<[0.,10.,2.,3.,4.,5.,6.,7.,8.,9.]> : tensor<10xf32> } : () -> ()

  %5 = "stablehlo.concatenate"(%4, %1) {dimension = 0 : i64} : (tensor<?xf32>,  tensor<?xf32>) -> tensor<?xf32>
  %6 = "quantfork.stats"(%5) {layerStats = dense<[-1.5726943, 4.6875381]> : tensor<2xf32>} : (tensor<?xf32>) -> tensor<?xf32>
  func.return  %3, %6 : tensor<?xf32>, tensor<?xf32>
}

// -----

  %2 = "tf.Mul"(%1, %value) {T = "tfdtype$DT_FLOAT"} : (tensor<256x8x7x16xf32>, tensor<16xf32>) -> tensor<256x8x7x16xf32>
  func.return %2 : tensor<256x8x7x16xf32>

// CHECK-DAG: %[[cst:.*]] = "tf.Const{{.*}} dense<8.000000e+00> : tensor<3x3x3x16xf32>
// CHECK-DAG: %[[cst_0:.*]] = "tf.Const{{.*}} dense<1.200000e+01> : tensor<16xf32>

func.func private @sum_reducer(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
  %0 = "tf.AddV2"(%arg0, %arg1) {device = ""} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
  func.return %0 : tensor<*xf32>
}

func.func @xla_func_0(%arg0: tensor<7xf32>, %arg1: tensor<f32>) -> tensor<10xf32> {
  %cst = "tf.Const"() {value = dense<0> : tensor<1x2xi32>} : () -> tensor<1x2xi32>

        module = ir.Module.parse(module_bytecode)
        self.assertIn('stablehlo.add %arg0, %cst : tensor<4xf32>', str(module))

  def test_tf_mlir_to_stablehlo(self):
    assembly = """
      module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
        func.func @main(%arg0 : tensor<?xf32>) -> tensor<?xf32> {
          %cst = "tf.Const"() {value = dense<1.0> : tensor<f32>} : () -> tensor<f32>

  %min = arith.constant dense<0.0> : tensor<16xf32>
  %max = arith.constant dense<15.0> : tensor<16xf32>
  %mini = "tf.Identity"(%min) : (tensor<16xf32>) -> tensor<16xf32>
  %maxi = "tf.Identity"(%max) : (tensor<16xf32>) -> tensor<16xf32>
  %fq = "tf.FakeQuantWithMinMaxVarsPerChannel"(%in, %mini, %maxi) {num_bits = 3, narrow_range = false} : (tensor<3x3x3x16xf32>, tensor<16xf32>, tensor<16xf32>) -> tensor<3x3x3x16xf32>

  %min = arith.constant dense<0.0> : tensor<16xf32>
  %max = arith.constant dense<255.0> : tensor<16xf32>
  %mini = "tf.Identity"(%min) : (tensor<16xf32>) -> tensor<16xf32>
  %maxi = "tf.Identity"(%max) : (tensor<16xf32>) -> tensor<16xf32>
  %fq = "tf.FakeQuantWithMinMaxVarsPerChannel"(%in, %mini, %maxi) {num_bits = 5, narrow_range = false} : (tensor<3x3x3x16xf32>, tensor<16xf32>, tensor<16xf32>) -> tensor<3x3x3x16xf32>

  %cst0 = mhlo.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
  // CHECK: %[[BROADCAST:.*]] = "mhlo.broadcast_in_dim"(%[[CONST]]) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<1x1x2x4xi32>
  %0 = "mhlo.broadcast_in_dim"(%cst0) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<1x1x2x4xi32>

module{
  func.func @main(%arg0: tensor<10xf32>) -> tensor<10xf32> {
    return %arg0 : tensor<10xf32>
  }
  func.func @func1(%arg0: tensor<10xf32>, %arg1: tensor<10xf32>) -> tensor<10xf32> {
    %0 = stablehlo.add %arg0, %arg1 : tensor<10xf32>
    %1 = stablehlo.add %arg0, %arg1 : tensor<10xf32>
    return %0 : tensor<10xf32>
  }
})mlir";