%max_iter = "tf.Const"() {__op_key = 1, value = dense<2> : tensor<i32>} : () -> tensor<i32>
  // CHECK: tf_mlrt.map_fn
  %result = "tf_mlrt.tf_map_fn"(%max_iter, %input2, %2) { operandSegmentSizes = array<i32: 1, 1, 1>, body_fn = @NopMapFnBody, num_tensor_list_or_flow_in = 1 : i32} : (tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>) -> tensor<i32>
  return %result : tensor<i32>
}

// CHECK-LABEL: @NopMapFnBody

      %outputs_22, %control_23 = tf_executor.island wraps "tf.Const"() {device = "/device:CPU:0", value = dense<3> : tensor<i32>} : () -> tensor<i32>
      %outputs_24, %control_25 = tf_executor.island wraps "tf.LessEqual"(%outputs_0, %outputs_22) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<*xi1>

    // core_tpu ops:
    %res_t = "tf.Identity"(%res_f) {_replication_info = "repl_info"} : (tensor<i32>) -> tensor<i32>

    // non_tpu_ops
    %res_n = "tf.Const"() {value = dense<2> : tensor<i32>} : () -> tensor<i32>

    return %res_n : tensor<i32>
  }
  func.func private @while_cond(%arg0: tensor<i32>) -> tensor<i1> {

  "tf.AssignVariableOp"(%arg0, %device_output) : (!tf_res_i32, tensor<i32>) -> ()
  "tf.SomeOp"(%device_output) : (tensor<i32>) -> ()
  func.return
}

// CHECK-LABEL: func @device_func_1
// CHECK-NOT: tf.aliasing_output
func.func @device_func_1(%arg0: tensor<i32>) -> tensor<i32> {
  func.return %arg0 : tensor<i32>
}

        %4:2 = "tf.TPUReplicatedOutput"(%1) : (tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<!tf_type.resource<tensor<i32>>>, tensor<!tf_type.resource<tensor<i32>>>)
        tf_executor.yield %arg0 : tensor<!tf_type.resource<tensor<i32>>>
      }
      tf_executor.fetch %outputs : tensor<!tf_type.resource<tensor<i32>>>
    }

func.func private @then(%x: tensor<i32>, %y: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
  func.return %x, %x : tensor<i32>, tensor<i32>
}

func.func private @else(%x: tensor<i32>, %y: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
  func.return %y, %y : tensor<i32>, tensor<i32>
}

// CHECK-LABEL: then/tfrt_dedup_results
// CHECK: return {{%.*}} : tensor<i32>

// CHECK-LABEL: else/tfrt_dedup_results

    %outputs_7, %control_7 = tf_executor.island wraps "tf.AddV2"(%outputs_5, %outputs) {_xla_compile_device_type = "TPU", _replication_info = "cluster1"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
    %outputs_8, %control_8 = tf_executor.island wraps "tf.AddV2"(%outputs_0, %outputs_3) {_xla_compile_device_type = "TPU"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
    tf_executor.fetch
  }

  %zps2 = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>
  %zps4 = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
  %one = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>

  // CHECK: %[[QUANT0:.*]] = mhlo.uniform_quantize %[[ARG:.*]] : (tensor<2x2xf32>) -> tensor<2x2x!quant.uniform<i8:f32, 1.000000e+00:2>>

  func.return %0 : tensor<f32>
}

// CHECK-LABEL: func @collective_reduce_v2_max_id
func.func @collective_reduce_v2_max_id(%input: tensor<f32>) -> tensor<f32> {
  %group_key = "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>
  %group_size = "tf.Const"() { value = dense<2> : tensor<i32> } : () -> tensor<i32>

  %one = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
  %elem_shape = "tf.Const"() {value = dense<[-1, 1]> : tensor<2xi32>} : () -> tensor<2xi32>
  %size = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
  // CHECK: TensorListReserve
  // CHECK: TensorListReserve
  %tl = "tf.TensorListReserve"(%elem_shape, %size) : (tensor<2xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<?x1xf32>>>