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

  AppendFunctionArguments(while_body, num_chains, chaining_data_type);
  AppendFunctionResults(while_body, num_chains, chaining_data_type);

  // Insert identity ops with control dep
  ChainResourceOps(while_body, chain_resource_to_ops_map,
                   resource_equivalence_classes, ops_connected_to_fetch,
                   num_old_outputs);
  // Modify all the while ops referencing the body function and the

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

    ```mlir
    func @cluster_with_loop() -> () {
      %0 = "tf.VarHandleOp"() ...
      "tf_device.cluster"() ( {
         %1 = "tf.While"(%0) {body = @while_body, cond = @while_cond}
         tf_device.return
      })
      return
    }
    func @while_body(%arg0: tensor<*x!tf_type.resource<tensor<f32>>>) {
     %constant = "tf.Const"() ...
     "tf.AssignVariableOp"(%arg0, %constant)
     return %arg0
    }

    %1 = "tf.While"(%0) {
      body = @while_body, cond = @while_cond, device = "", is_stateless = false}
         : (tensor<*x!tf_type.resource<tensor<f32>>>) -> (tensor<*x!tf_type.resource<tensor<f32>>>)
    tf_device.return
  }) {cluster_attr = "cluster_attr"} : () -> ()
  func.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 ...
   }