func.func @convert_cluster_func(%arg0: tensor<i32>) -> tensor<i32> {
    // CHECK: "tf.XlaLaunch"(%arg0) <{function = @func, operandSegmentSizes = array<i32: 0, 1, 0>}> : (tensor<i32>) -> tensor<i32>
    %0 = "tf_device.cluster_func"(%arg0) {func = @func} : (tensor<i32>) -> tensor<i32>
    func.return %0 : tensor<i32>
  }

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

// -----

  %1 = "tf.ReadVariableOp"(%arg) {device = "/device:CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
  %2 = "tf.Add"(%1, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
  "tf.AssignVariableOp"(%arg, %2) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
  func.return %arg : tensor<!tf_type.resource<tensor<i32>>>
}

func.func private @"random/while_cond/TfMlrtAsyncWhilePredicate"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {

  %0 = "tf.AddV2"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
  return %0: tensor<i32>
}

// CHECK-LABEL: func.func private @"random/while_body/TfMlrtAsyncWhileBody"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {

  "tf.AssignVariableOp"(%arg1, %0#1) : (tensor<*x!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
  func.return
}

// CHECK-LABEL: func @multiple_write_func
// CHECK-SAME: ({{%.*}}: tensor<i32>) -> (tensor<i32>, tensor<i32>)
func.func @multiple_write_func(%arg0: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
  func.return %arg0, %arg0 : tensor<i32>, tensor<i32>
}

    "tfl.yield"(%1#0, %1#1) : (tensor<*xi32>, tensor<*xf32>) -> ()
  }) : (tensor<i32>, tensor<1xf32>) -> (tensor<i32>, tensor<1xf32>) loc("WhileOp")
  // CHECK: (tensor<i32>, tensor<1xf32>, tensor<i32>) ->
  // CHECK-SAME: (tensor<i32>, tensor<1xf32>, tensor<i32>)
  func.return %0#1 : tensor<1xf32>
}

func.func private @WhileOp_cond(%arg0: tensor<*xi32>, %arg1: tensor<*xf32>, %arg2: tensor<i32>) -> tensor<i1> {

  func.return %0 : tensor<i32>
}

func.func @inner_stateful_pcall_func(%arg0: tensor<i32>) -> tensor<i32> {
  %0 = "tf.StatefulPartitionedCall"(%arg0) {_xla_compile_device_type = "CPU", config = "", config_proto = "", device = "/device:CPU:0", executor_type = "", f = @func} : (tensor<i32>) -> (tensor<i32>)
  func.return %0 : tensor<i32>
}

// CHECK-SAME:    (%[[ARG0:.*]]: tensor<i32>)
func.func @identity(%arg0: tensor<i32>) -> tensor<i32> {
  // CHECK-NOT: "tf.Identity"
  %0 = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<i32>
  // CHECK: return %[[ARG0]]
  func.return %0 : tensor<i32>
}

// CHECK-LABEL: func @identity_terminator
// CHECK-SAME:    (%[[ARG0:.*]]: tensor<i32>)
func.func @identity_terminator(%arg0: tensor<i32>) -> (tensor<*xi32>, tensor<i32>) {

  %outputs_22 =  "tf.Const"() {device = "", value = dense<16> : tensor<i32>} : () -> tensor<i32>
  %outputs_24 =  "tf.Const"() {device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
  %outputs_26 =  "tf.Range"(%outputs_24, %outputs_22, %outputs_20) {device = ""} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<*xi32>
  // CHECK: [[read_only_tensor_list:%.*]] = "tf.TensorListFromTensor"

  %2 = "tf.AddV2"(%y, %1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
  func.return %1, %2 : tensor<i32>, tensor<i32>
}

func.func @no_side_effect_then_1(%x: tensor<i32>, %y: tensor<i32>) -> tensor<i32> {
  %0 = "tf.AddV2"(%x, %y) : (tensor<i32>, tensor<i32>) -> tensor<i32>
  func.return %0 : tensor<i32>
}

func.func @no_side_effect_else_1(%x: tensor<i32>, %y: tensor<i32>) -> tensor<i32> {

      %replicated_output:2 = "tf.TPUReplicatedOutput(%identity) : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
      return %replicated_output#0, %replicated_output#1 : tensor<i32>, tensor<i32>
    }
    ```

    will be transformed into:

    ```mlir
    func @tpu_computation(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, tensor<i32>) {