erase_indices.set(i);
    } else if (results[i].getDefiningOp()) {
      FindConsumedOp(results[i].getDefiningOp(), consumed_ops);
    }
  }
  // Remove elements and operations in while_body that are not indispensable.
  yield_op.eraseOperands(erase_indices);
  // Remove ops from bottom to top.
  for (Operation &op :
       llvm::make_early_inc_range(reverse(body.getOperations())))

  // CHECK-NOT: tf.TensorListPopBack
  // CHECK: return
  func.return
}
// CHECK: func @while_body(%[[BARG0:.*]]: tensor<10xf32>, %[[BARG1:.*]]: tensor<i32>, %[[BARG2:.*]]: tensor<1xi32>)
func.func @while_body(%arg0: tensor<!tf_type.variant<tensor<f32>>>, %arg1: tensor<i32>) -> (tensor<!tf_type.variant<tensor<f32>>>, tensor<i32>) {

  // functions and avoid converting ConstOps inside initializer functions.
  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    // Do not unfreeze constants under these functions.
    if (func_op.getSymName().contains("while_body")) continue;
    if (func_op.getSymName().contains("while_cond")) continue;
    absl::c_copy_if(func_op.getOps<TF::ConstOp>(),
                    std::back_inserter(target_const_ops),

  func.return %0 : tensor<f32>
}

// -----

// Tests `mhlo.while` with only the `body` region populated with TF/XLA
// communication ops.

// CHECK-LABEL: func @while_body
// CHECK-SAME:  ([[ARG0:%.*]]: tensor<f32>)
func.func @while_body(%arg0: tensor<f32>) -> tensor<f32> {
  // CHECK: [[INIT_TOKEN:%.*]] = mhlo.create_token

    // expected-remark@above {{Predecessors: {4}}}
  }
  func.return
  // expected-remark@above {{ID: 7}}
  // expected-remark@above {{Sinks: {6}}}
}

// CHECK-LABEL: func @while_body
func.func @while_body(%arg0: tensor<i1>) -> tensor<i1> {
  // expected-remark@above {{ID: 5}}
  %graph = tf_executor.graph {
  // expected-remark@above {{ID: 3}}
    %island:2 = tf_executor.island {

    variable ops, thus allowing them to run in parallel across iterations.

    Before:

    ```mlir
    !tf_res = type tensor<!tf_type.resource<tensor<f32>>>
    func @while_body(%arg0: !tf_res, %arg1: !tf_res, %arg2: tensor<f32>, %arg3: tensor<f32>) -> (!tf_res, !tf_res, tensor<f32>, tensor<f32>) {
      %graph:4 = tf_executor.graph {

variable ops, thus allowing them to run in parallel across iterations.

Before:

```mlir
!tf_res = type tensor<!tf_type.resource<tensor<f32>>>
func @while_body(%arg0: !tf_res, %arg1: !tf_res, %arg2: tensor<f32>, %arg3: tensor<f32>) -> (!tf_res, !tf_res, tensor<f32>, tensor<f32>) {
  %graph:4 = tf_executor.graph {

  // Use this to clone an op and lift it outside its parent function. The
  // original while body is unchanged. Example:
  // Original:
  //    %x = tf.while(%a, %b)
  //    ...
  //    while_body:
  //       call(f=@sc_fw, %arg0, %arg1)
  // Lifted:
  //    call(f=@sc_fw, %a, %b)
  //    %x = tf.while(%a, %b)
  //    ...
  func::FuncOp orig_parent_func =