tf_device.return %arg0, %arg1, %arg2, %arg3, %arg4 : tensor<i1>, tensor<i32>, tensor<i64>, tensor<f32>, tensor<f64>
  // CHECK-NEXT: }, {
  }, {
    // CHECK-NEXT: tf_device.return
    tf_device.return
  // CHECK-NEXT: }) : () -> (tensor<i1>, tensor<i64>, tensor<f64>)
  }) : () -> (tensor<i1>, tensor<i32>, tensor<i64>, tensor<f32>, tensor<f64>)

  %0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
    quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
  } : (tensor<1xf32>, tensor<f32>, tensor<i32>) -> tensor<1x!tf_type.qint8>
  %1 = "tf.UniformDequantize"(%0, %scales, %zps) {
    quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
  } : (tensor<1x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1xf32>

    %outputs_0, %control_1 = tf_executor.island wraps "tf.LegacyCall"(%outputs, %outputs) {_disable_call_shape_inference = true, f = @test_func_name0} : (tensor<!tf_type.resource<tensor<i64>>>, tensor<!tf_type.resource<tensor<i64>>>) -> tensor<*x!tf_type.resource> loc("called")

// CHECK-SAME:       {program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []
// CHECK-SAME:       (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
// CHECK-NEXT:    %1 = "tf.Identity"(%arg1) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>
// CHECK-NEXT:    %2 = "tf.IfrtCall"(%1, %arg0)
// CHECK-SAME:       {program_id = [[PROGRAM_ID]] : i64, variable_arg_indices = []

  %filter = mhlo.constant dense<[[[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]]]> : tensor<1x1x3x2xf32>

  // CHECK-NOT: corert_tensorhandle_to_fallback_tensor

  %0 = tfrt_fallback_async.fallback_tensor_to_corert_tensorhandle %arg {_tfrt_cost = 1 : i64, device = "/CPU:0"} : (!tfrt_fallback.tf_tensor) -> (!corert.tensorhandle)
  %1 = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %0 {_tfrt_cost = 1 : i64, device = "/CPU:0"} : (!corert.tensorhandle) -> (!tfrt_fallback.tf_tensor)

  // CHECK: tfrt.return [[arg]]

func.func @serving_default(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
  %0 = "tf.BatchFunction"(%arg0, %arg1) {f = @batched_func, num_batch_threads = 1 : i64, max_batch_size = 2 : i64, batch_timeout_micros = 10000 : i64, operandSegmentSizes = array<i32: 1, 1>} : (tensor<1xf32>, tensor<1xf32>) -> (tensor<1xf32>)
  return %0 : tensor<1xf32>
}
// The contents of `@serving_default` should have been inlined to `@batch_func`.

    %3:2 = tf_executor.island wraps "tf.While"(%1#0) {name = "B", body = @while_body_with_wrong_cluster_attr, cond = @while_cond_with_wrong_cluster_attr, is_stateless = false, parallel_iterations = 10 : i64} : (tensor<i32>) -> tensor<i32>

    %0 = "tf.Equal"(%arg0, %cst) {_tpu_replicate = "cluster", device = "", incompatible_shape_error = true} : (tensor<*xi64>, tensor<i64>) -> tensor<*xi1>

  tf_executor.graph {