%3 = "tf.Mul"(%arg0, %0) : (tensor<16xf32>, tensor<f32>) -> tensor<16xf32>
        %4 = "tf.Mul"(%3, %0) : (tensor<16xf32>, tensor<f32>) -> tensor<16xf32>
        %5 = "tf.Mul"(%4, %1) : (tensor<16xf32>, tensor<f32>) -> tensor<16xf32>
        %6 = "tf.Mul"(%5, %2) : (tensor<16xf32>, tensor<f32>) -> tensor<16xf32>
        tf_device.return %6 : tensor<16xf32>
      }) {} : () -> tensor<16xf32>
      tf_executor.yield %3 : tensor<16xf32>
    }

module {
  // CHECK-LABEL: func @main
  // CHECK-SAME: %[[ARG_0:.*]]: tensor<10xf32>, %[[ARG_1:.*]]: tensor<10xf32>, %[[ARG_2:.*]]: tensor<6xi32>
  func.func @main(
      %arg0: tensor<10xf32>, %arg1: tensor<10xf32>, %arg2: tensor<6xi32>
    ) -> (tensor<10xf32>, tensor<6xi32>) {
    // CHECK: %[[CAST_0:.*]] = "tf.Cast"(%[[ARG_0]]) <{Truncate = false}> : (tensor<10xf32>) -> tensor<10xbf16>

  %0 = "quantfork.stats"(%arg0) {bitsNum = 8 : i64, layerStats = dense<[-0.999415695, 0.99998933]> : tensor<2xf32>, narrowRange = false} : (tensor<?x3xf32>) -> tensor<?x3xf32>
  // CHECK: %[[dot:.*]] = stablehlo.dot %[[dq2]], %[[dq1]]
  %1 = stablehlo.dot %0, %cst : (tensor<?x3xf32>, tensor<3x2xf32>) -> tensor<?x2xf32>
  // CHECK: %[[q3:.*]] = "quantfork.qcast"(%[[dot]])
  // CHECK-SAME: quant.uniform<i4:f32, 0.61633392969767253:-2>>

// CHECK: func @while_cond(%arg0: tensor<256x256xi32>)
// CHECK: func private @while_body_0(%arg0: tensor<128xi32>)
// CHECK: func private @while_cond_1(%arg0: tensor<128xi32>)
func.func @while_main(%arg0: tensor<256x256xi32>, %arg1: tensor<128xi32>) -> (tensor<256x256xi32>, tensor<128xi32>) {

// CHECK-NEXT:   "tfl.yield"(%[[MY_CUSTOM]]) : (tensor<4xf32>) -> ()
// CHECK-NEXT: }) {fused_activation_function = "RELU", int_attr = 2 : i32} : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
// CHECK-NEXT: %[[CUSTOM_2:.*]] = "tfl.custom_tf"(%[[MUL]], %[[CST]]) ({
// CHECK-NEXT: ^bb0(%arg1: tensor<4xf32>, %arg2: tensor<4xf32>):

  %4 = "tfr.cast"(%concat) : (!tfr.tensor) -> tensor<3xf32>
  func.return %4 : tensor<3xf32>

// CHECK: %[[concat:.*]] = "tf.RiscConcat"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<*xf32>
// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[concat]]) <{shape = #tf_type.shape<3>}> : (tensor<*xf32>) -> tensor<3xf32>
// CHECK: return %[[es]] : tensor<3xf32>
}

// CHECK-LABEL: decompose_tf_split

  %0 = "tf.TensorListFromTensor"(%arg0, %arg1) : (tensor<3x10xf32>, tensor<1xi32>) -> tensor<!tf_type.variant<tensor<10xf32>>>
  %1 = "tf.TensorListSetItem"(%0, %arg2, %arg3) : (tensor<!tf_type.variant<tensor<10xf32>>>, tensor<i32>, tensor<10xf32>) -> tensor<!tf_type.variant<tensor<10xf32>>>

  # Note that semantically in Python, multiple return values are equivalent
  # to returning a tuple/list.
  #
  # CHECK:      func {{@[a-zA-Z_0-9]+}}() -> (
  # CHECK-SAME:   tensor<1xf32> {tf_saved_model.index_path = [0]},
  # CHECK-SAME:   tensor<2xf32> {tf_saved_model.index_path = [1]})
  # CHECK-SAME: attributes {{.*}} tf_saved_model.exported_names = ["f0001_multiple_results_no_punctuation"]
  @tf.function(input_signature=[])

func.func private @simple_callee_with_devices() -> (tensor<2xi32>, tensor<2xf32>)  {
  %cst = "tf.Const"() { value = dense<2> : tensor<2xi32> } : () -> tensor<2xi32>
  %cst_gpu = "tf.Const"() { value = dense<2.> : tensor<2xf32>, device = "/GPU:0" } : () -> tensor<2xf32>
  func.return %cst, %cst_gpu : tensor<2xi32>, tensor<2xf32>
}
// CHECK-LABEL: func @inline_simple_with_devices(

    %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> } : () -> ()