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);

                      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>) {

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

// CHECK: func private @while_body
// CHECK-LABEL: testWhileRegionTrivialCasts
func.func private @while_cond(%arg0 : tensor<4xf32>, %arg1 : tensor<i32>) -> tensor<i1>
func.func private @while_body(%arg0 : tensor<4xf32>, %arg1 : tensor<i32>) -> (tensor<4xf32>, tensor<i32>)
func.func @testWhileRegionTrivialCasts(%arg0 : tensor<*xf32>, %arg1 : tensor<i32>) -> tensor<*xf32> {
  // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) <{body = @while_body, cond = @while_cond

  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

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

    return %cst_0 : tensor<8xf32>
  }
}


// -----

// Tests the case when there are functions called from the main function such as while_body/while_cond.

module attributes {tf_saved_model.semantics} {