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

  // CHECK-NEXT: return
  %const = "tf.Const"() {device = "/CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
  %x = "tf.AddV2"(%const, %const) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %r = "tf.AddV2"(%arg, %x) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
  func.return %r : 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>
}

func.func @target_function() -> tensor<i32> {
  %0 = "tf_device.cluster"() ({
    // CHECK: tf.UnsupportedOp
    %1 = "tf.UnsupportedOp"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
    // CHECK: tf.Identity
    %2 = "tf.Identity"(%1) : (tensor<i32>) -> tensor<i32>
    // CHECK-NOT: tf.Identity
    %3 = "tf.Identity"(%2) : (tensor<i32>) -> tensor<i32>
    %4 = "tf.Identity"(%3) : (tensor<i32>) -> tensor<i32>
    // CHECK: tf_device.return

  func.return %0 : tensor<? x i32>
}

// CHECK-LABEL: testPow
func.func @testPow(tensor<? x i32>, tensor<? x i32>) -> tensor<? x i32> {
^bb0(%arg0: tensor<? x i32>, %arg1: tensor<? x i32>):
  // CHECK: tfl.pow %arg0, %arg1
  %0 = tfl.pow %arg0, %arg1 : tensor<? x i32>
  func.return %0#0 : tensor<? x i32>
}

// CHECK-LABEL: testAtan2