%0 = "tfl.conv_2d"(%arg0, %arg1, %arg2) {dilation_h_factor = 2 : i32, dilation_w_factor = 2 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<1x112x80x128xf32>, tensor<128x3x3x128xf32>,...

                      %n: i32 {tfr.name="N"},
                      %axis: i32 {tfr.name="axis"}) -> !tfr.tensor {
  %index = arith.constant 0 : index
  %cst = arith.constant 1 : i32
  %eq = arith.cmpi eq, %n, %cst : i32
  %v1 = tfr.get_element %values[%index] : (!tfr.tensor_list, index) -> !tfr.tensor
  %temp = tfr.call @tf__expand_dims(%v1, %axis) : (!tfr.tensor, i32) -> !tfr.tensor
  %res = scf.if %eq -> !tfr.tensor {

  %1 = "tf.Const"() {value = dense<5> : tensor<i32>} : () -> tensor<i32>
  %2 = "tf.Add"(%0, %1) : (tensor<i32>, tensor<i32>) -> (tensor<i32>)
  func.return %2 : tensor<i32>
}

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

// -----

// CHECK: "tf.While"
// CHECK-SAME: (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant>) -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant>)

func.func @otherVariantWhileBody(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<!tf_type.variant>) -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant>) {
  %0 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>

    %one:2 = tf_executor.island wraps "tf.Const"() { value = dense<1> : tensor<i32> } : () -> tensor<i32>
    %add2:2 = tf_executor.island wraps "tf.Add"(%arg0, %one#0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
    tf_executor.fetch %add2#0, %arg1, %stack_push2#1 : tensor<i32>, tensor<f32>, !tf_executor.control
  }
  func.return %graph#0, %graph#1 : tensor<i32>, tensor<f32>
}
// CHECK-LABEL:   func @unique_resource_chain_while_body

  "func.func"() ({
  ^bb0(%arg0: tensor<i32>):
    "func.return"(%arg0) : (tensor<i32>) -> ()
  }) {sym_name = "main_while_body_4225150", sym_visibility = "private", tf._input_shapes = [#tf_type.shape<>], tf.signature.is_stateful, function_type = (tensor<i32>) -> (tensor<i32>)} : () -> ()
  "func.func"() ({
  ^bb0(%arg0: tensor<i32>):

     %2 = builtin.unrealized_conversion_cast to tensor<i32>
     %3 = builtin.unrealized_conversion_cast to tensor<i32>
     // CHECK: TPUCopyWithDynamicShape{{.*}}device = "foobar"
     %4, %5 = "tf.TPUCopyWithDynamicShape"(%0, %1, %2, %3) {operandSegmentSizes = array<i32: 2, 2>} : (tensor<2048xi32>, tensor<2048xi32>, tensor<i32>, tensor<i32>) -> (tensor<2048xi32>, tensor<2048xi32>)

  %cst = "tf.Const"() {device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
  %cst_0 = "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
  %0 = "tf.AddV2"(%cst, %cst_0) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %1 = "tf.Max"(%arg0, %0) {device = ""} : (tensor<1x2x3xf32>, tensor<i32>) -> tensor<1x3xf32>

  return %1: tensor<1x3xf32>
}

func.func @while_cond(%arg0: tensor<i32>) -> tensor<i1> {
  %0 = "tf.Const"() {value = dense<9> : tensor<i32>} : () -> tensor<i32>
  %1 = "tf.Less"(%arg0, %0) {} : (tensor<i32>, tensor<i32>) -> tensor<i1>
  func.return %1 : tensor<i1>
}

func.func @while_body(%arg0: tensor<i32>) -> tensor<i32> {
  %1 = "tf.AddV2"(%arg0, %arg0) {} : (tensor<i32>, tensor<i32>) -> tensor<i32>
  func.return %1 : tensor<i32>
}

  func.return %0 : tensor<i32>
}

// -----

func.func @Callee(%arg0: tensor<i32>) -> tensor<i32> {
  "tf.XlaSendToHost"(%arg0) {key = "send_key_call3"} : (tensor<i32>) -> ()
  %1 = "tf.XlaRecvFromHost"() {key = "recv_key_call3", shape = #tf_type.shape<>} : () -> tensor<i32>
  func.return %1 : tensor<i32>
}