%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

    // 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> {

  %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>>

  %0 = "tfl.pseudo_const"() {value = dense<[2, 2, 81]> : tensor<3xi32>} : () -> tensor<3xi32>
  %1 = "tfl.pseudo_const"() {value = dense<-1> : tensor<i32>} : () -> tensor<i32>
  %2:3 = "tfl.split_v"(%arg0, %0, %1) {num_splits = 3 : i32} : (tensor<?x13x13x85xf32>, tensor<3xi32>, tensor<i32>) -> (tensor<?x13x13x2xf32>, tensor<?x13x13x2xf32>, tensor<?x13x13x81xf32>)

  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>

  func.return %0 : tensor<i32>
}

// -----

func.func @unsupported_op_soft_placement_false() -> tensor<i32> {
  %0 = "tf_device.cluster"() ({
    // expected-error@+1 {{'tf.UnsupportedOp' op isn't compilable for TPU device}}
    %1 = "tf.UnsupportedOp"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
    %2 = "tf.Identity"(%1) : (tensor<i32>) -> tensor<i32>
    tf_device.return %2 : tensor<i32>

  {
    %cst = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
    %0 = "tf.AddV2"(%arg2, %cst) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
    %1 = "tf.Identity"(%0) {device = ""} : (tensor<i32>) -> tensor<i32>

  %stride = "tf.Add"(%arg2, %cst) : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %cst1 = arith.constant dense<44> : tensor<i32>
  %limit = "tf.Add"(%arg2, %cst1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
  // CHECK: [[Result:%.*]]:4 = "tf.While"(%arg0, %arg1, %{{.+}}, %{{.+}} <{body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond
  %0:2 = "tf.WhileRegion"(%arg0, %arg1) (
    {
      ^bb0(%carg0: tensor<*xf32>, %carg1: tensor<i32>):

    %cst = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
    %11:4 = "tf.Split"(%cst, %0) {num_split = 4 : i32} : (tensor<i32>, tensor<128x1024xf32>) -> (tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>, tensor<32x1024xf32>)
    %cst_0 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>

    ^bb0 (%carg0: tensor<!tf_type.resource>, %carg1: tensor<i32>):
    %pred = "tf._SomeOp"(%carg1) : (tensor<i32>) -> tensor<i1>
    "tf.Yield"(%pred) : (tensor<i1>) -> ()
  }, {
    ^bb0 (%carg0: tensor<!tf_type.resource>, %carg1: tensor<i32>):
    %const1 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
    %sub = "tf.Sub"(%carg1, %const1) : (tensor<i32>, tensor<i32>) -> tensor<i32>