%0 = "tf.While"(%arg0) {body = @while_body, cond = @while_cond, device = "", is_stateless = true} : (tensor<256x256xi32>) -> (tensor<256x256xi32>)
  %1 = "tf.While"(%arg1) {body = @while_body, cond = @while_cond, device = "", is_stateless = true} : (tensor<128xi32>) -> (tensor<128xi32>)
  func.return %0, %1: tensor<256x256xi32>, tensor<128xi32>
}

func.func @while_body(%arg0: tensor<*xi32>) -> (tensor<*xi32>) {

  func.return %1 : tensor<i1>
}

// The @while_body function is a Multi-hosts function which contains three
// parts:
//   - A local subgraph which contains both local ops and remote_run kernels to
//     trigger remote subgraphs
//   - Two remote subgraph which contains remote ops on worker:1 and worker:2.
// CHECK: func @while_body(%[[ARG_0:.*]]: tensor<i32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"})

func.func @while_cond(%arg0: tensor<*xf32>) -> tensor<i1> {
  %0 = "tf.Const"() {value = dense<true> : tensor<i1>} : () -> tensor<i1>
  func.return %0 : tensor<i1>
}

func.func @while_body(%arg0: tensor<*xf32>) -> (tensor<*xf32>) {
  %0 = "tf.SomeOp"(%arg0) : (tensor<*xf32>) -> tensor<*xf32>
  func.return %0 : tensor<*xf32>
}

// Test that -tf-drop-while-shape-invariant-in-device-cluster pass does not drop

  %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:   yield
// CHECK: ^bb0([[ARGS]]):
// CHECK:   call @while_body
// CHECK:   yield
// CHECK: while_body
// CHECK: while_cond

// INLINE: tfl.while
// INLINE: ^bb0([[ARGS:.*]]):
// INLINE:   tf.Less
// INLINE:   yield
// INLINE: ^bb0([[ARGS]]):
// INLINE:   %cst_2 = arith.constant
// INLINE:   yield
// INLINE-NOT: while_body
// INLINE-NOT: while_cond

// CANON-LABEL:   func @while_main

    // generated module.
    auto& child_node = nodes.at(child.node_id());
    auto child_name = child.local_name().c_str();

    if (child_node.kind_case() == tensorflow::SavedObject::kFunction) {
      Function tf_function;
      for (const std::string& fn_name :
           child_node.function().concrete_functions()) {
        // Setup input signature.
        //

      %control = tf_executor.island {
        %0 = "tf.While"(%arg0) <{body = @while_body, cond = @while_cond, is_stateless = false, shape_invariant}> : (tensor<*x!tf_type.resource>) -> tensor<!tf_type.resource<tensor<i32>>>
        tf_executor.yield
      }
      tf_executor.fetch %control : !tf_executor.control
    }
    return
  }

      %1, %c0 = tf_executor.island wraps "tf.Const"() {value = dense<1> : tensor<i1>} : () -> tensor<i1>
      tf_executor.fetch %1 : tensor<i1>
    }
    return %0 : tensor<i1>
  }

  func.func private @while_body(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
    tf_executor.graph {

    input.

#### Pseudo MLIR code before and after RewritePass

```mlir
// Example: original program:
//
module {
   func @while_body {
     %input = "tf.IteratorGetNext"(...) {device = "/CPU:0"}:
              -> tensor<2x224x224x3xf32>
     %device_launch = "tf_device.launch_func"(%input,...) {func = @_func,...)
     return ...
   }

# This check that we don't error out when importing GraphDef containing
# functions with arg name that are the same as the graph input name

# CHECK: func @main(%arg0: tensor<{{.*}}i32>) -> tensor<{{.*}}i32>
# CHECK: func private @while_body
# CHECK: func private @while_cond

node {
  name: "input"
  op: "Placeholder"
  attr {
    key: "_user_specified_name"
    value {
      s: "input"
    }
  }
  attr {
    key: "dtype"