EXPECT_EQ(input_node->name(), "arg0");

  // Check node "while" input and output edges.
  const Node *while_node = node_name_index.at("while");
  ASSERT_NE(while_node, nullptr);
  TF_CHECK_OK(while_node->input_node(0, &input_node));
  EXPECT_EQ(input_node->name(), "arg1");
  TF_CHECK_OK(while_node->input_node(1, &input_node));
  EXPECT_EQ(input_node->name(), "arg0");
  const Node *ret2_node = node_name_index.at("ret2");

  is_compilable &= ExtractNodeDefAndCheckCompilability(
      while_node, "cond", "while_cond", encapsulating_function, lib_runtime,
      stack_trace, uncompilable_nodes);

  if (!uncompilable_nodes && !is_compilable) return is_compilable;

  is_compilable &= ExtractNodeDefAndCheckCompilability(
      while_node, "body", "while_body", encapsulating_function, lib_runtime,
      stack_trace, uncompilable_nodes);

    TF_RETURN_IF_ERROR(GraphToFunctionDef(*g, new_name, &replace_fdef));
    TF_RETURN_IF_ERROR(fld->AddFunctionDef(replace_fdef));
    loop_cond_func->set_name(new_name);
    while_node->ClearAttr("cond");
    while_node->AddAttr("cond", *loop_cond_func);
  }

  return absl::OkStatus();
}

// Rewrites while loop cond function for host.
Status RewriteHostWhileLoopCond(

                      NameAttrList* encapsulating_function,
                      UncompilableNodesMap* uncompilable_nodes) const;
  bool IsCompilableWhile(const Node& while_node,
                         FunctionLibraryRuntime* lib_runtime,
                         std::vector<StackFrameView>* stack_trace,
                         NameAttrList* encapsulating_function,

    // CHECK: {{.*tf.While.*body = @while_body.* cond = @while_cond.* parallel_iterations = 3}}
    %0 = "tf.While"(%cst) {body = @while_body, cond = @while_cond, is_stateless = false, parallel_iterations = 3} : (tensor<i32>) -> (tensor<i32>)
    return
  }
  func.func private @while_body(%arg0: tensor<i32>) -> (tensor<i32>) {
    // metadata ops

    %0 = "tf.While"(%cst_main) {body = @while_body, cond = @while_cond, is_stateless = false} : (tensor<i32>) -> (tensor<i32>)
    // expected-error @+1 {{'tf.While' op multiple users of function.}}
    %1 = "tf.While"(%cst_main) {body = @while_body, cond = @while_cond, is_stateless = false} : (tensor<i32>) -> (tensor<i32>)
    return
  }
  func.func private @while_body(%arg0: tensor<i32>) -> (tensor<i32>) {

  %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"})

// CHECK-NEXT:    %0 = "tf.Less"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK-NEXT:    return %0 : tensor<i1>

// CHECK-LABEL: func.func private @"map/while_body"

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