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

  %6 = "tf.Identity"(%1) : (tensor<!tf_type.int32ref>) -> tensor<i32>

  func.return %6 : tensor<i32>
}

// CHECK-LABEL: @inference
func.func @inference() -> (tensor<i32>, tensor<i32>, tensor<i32>) {
  // CHECK-NOT: tf.VariableV2

  // CHECK: [[handle:%.*]] = "tf.VarHandleOp"
  // CHECK-SAME: shared_name = "x"
  // CHECK: "tf.ReadVariableOp"([[handle]])

  %0:3 = "tfl.while"(%cst_0, %cst_1, %cst_2) (
    // cond
    {
    ^bb0(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<i32>):
      %cst_3 = "arith.constant" () {value = dense<3> : tensor<i32>} : () -> tensor<i32>
      %1 = "tfl.less"(%arg2, %cst_3) : (tensor<i32>, tensor<i32>) -> tensor<i1>
      "tfl.yield"(%1) : (tensor<i1>) -> ()
    },
    // body
    {

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

// CHECK: while_1([[ARG0:%[^ :]*]]: tensor<i32>, [[ARG1:%[^ :]*]]: tensor<1xf32>)
func.func @while_1(%arg0: tensor<i32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
  // CHECK: [[CST:%[^ ]*]] = arith.constant dense<1> : tensor<i32>
  // CHECK: "tf.WhileRegion"([[ARG0]], [[ARG1]])
  // CHECK: (tensor<i32>, tensor<1xf32>) -> (tensor<i32>, tensor<1xf32>)
  %0:2 = "tf.WhileRegion"(%arg0, %arg1) (
    // cond
    {

func.func private @batched_function(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32>
  attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
  // CHECK: tf.Const
  %1 = "tf.AddV2"(%arg0, %arg1) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %2 = "tf.Identity"(%1) {device = "/device:CPU:0"} : (tensor<i32>) -> tensor<i32>
  func.return %2 : tensor<i32>
}

// CHECK-LABEL: func @main

    %0 = "tf.Const"() {value = dense<100> : tensor<i32>} : () -> tensor<i32>
    %1 = "tf.WhileRegion"(%0) ({
       // Condition region
       ^bb0(%carg0: tensor<i32>):
          %c0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
          %c1 = "tf.GreaterEqual"(%carg0, %0) {T = i32, device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i1>
          "tf._UnknownOp1_"(%arg1) : (!tf_res_f32) -> ()