ASSERT_NE(if_node, nullptr);
  const Node *input_node;
  TF_CHECK_OK(if_node->input_node(1, &input_node));
  EXPECT_EQ(input_node->name(), "arg1");
  TF_CHECK_OK(if_node->input_node(2, &input_node));
  EXPECT_EQ(input_node->name(), "arg0");
  const Node *ret0_node = node_name_index.at("ret0");
  ASSERT_NE(ret0_node, nullptr);
  TF_CHECK_OK(ret0_node->input_node(0, &input_node));
  EXPECT_EQ(input_node->name(), "if");

              shape_refiner->GetContext(n);
          for (int i = 0; i < n->num_inputs(); i++) {
            const Node* input_node;
            if (n->input_node(i, &input_node).ok()) {
              auto shapes_and_types = context->input_handle_shapes_and_types(i);
              if (shapes_and_types) {
                context->set_output_handle_shapes_and_types(0,

  EXPECT_EQ(identity_node->attrs().FindByString("T")->type(), DT_FLOAT);

  EXPECT_EQ(identity_node->num_inputs(), 1);
  const Node* identity_input_node = nullptr;
  TF_EXPECT_OK(identity_node->input_node(0, &identity_input_node));
  EXPECT_EQ(identity_input_node->name(), "_arg0");

  const Node* arg_node = node_name_index.at("_arg0");
  EXPECT_EQ(arg_node->op_def().name(), "_Arg");

      send_recv_nodes.push_back(n);
    }
  }
  EXPECT_EQ(num_send_from_host, 1);
  EXPECT_EQ(num_recv_at_host, 1);
  for (Node *n : send_recv_nodes) {
    Node *input_node;
    TF_CHECK_OK(n->input_node(n->num_inputs() - 1, &input_node));
    EXPECT_EQ(input_node, key_placeholder);

    bool has_control_edge_to_sequencer = false;
    for (const Edge *e : n->out_edges()) {
      if (e->IsControlEdge() && e->dst() == sequencer) {

  NodeBuilder while_builder(kFunctionalWhileNodeName, "While",
                            builder.opts().op_registry());
  while_builder.Input({input_node, input_node})
      .Attr("cond", compilable)
      .Attr("body", uncompilable);
  builder.opts().FinalizeBuilder(&while_builder);

  GraphDef graph_def;
  TF_EXPECT_OK(builder.ToGraphDef(&graph_def));
  std::unique_ptr<Graph> graph(new Graph(flib_def_.get()));

  // Create a list of all control dependencies to be copied when possibly
  // replacing nodes related to slice_size.
  Node* old_size;
  std::vector<const Edge*> old_size_ctrl_deps;
  TF_RETURN_IF_ERROR(slice->input_node(2, &old_size));
  absl::c_copy_if(old_size->in_edges(), std::back_inserter(old_size_ctrl_deps),
                  [](const Edge* e) { return e->IsControlEdge(); });

  Output slice_size;

  for (const auto* input_edge : in_edges) {
    const Node& input_node = *input_edge->src();
    if (input_node.IsSource()) {
      if (in_edges.size() != 1) {
        return errors::FailedPrecondition(
            "The node has other inputs besides the _Source node");
      }
      // We don't import the _SOURCE node.
      continue;
    }
    if (input_node.IsArg() && input_edge->IsControlEdge()) {

  }
  if (!ret_node) {
    return errors::Internal("No _Retval node for loop cond function ",
                            loop_cond_func->name());
  }
  Node* loop_cond;
  TF_RETURN_IF_ERROR(ret_node->input_node(0, &loop_cond));

  // Build the XlaSendToHost node.
  NodeDefBuilder send_loop_cond_builder(
      absl::StrCat("send_oc_while_cond_", while_node->name()), "XlaSendToHost");

    return false;
  }
  const Node* switch_node = (*it)->src();

  // Check if the Switch is driven by LoopCond.
  const Node* maybe_loop_cond;
  TF_RETURN_IF_ERROR(switch_node->input_node(1, &maybe_loop_cond));
  if (!maybe_loop_cond->IsLoopCond()) {
    return false;
  }

  // Check if the Identity is driving any const nodes through a control edge.
  bool driving_any_consts =

                                    " into function '", fn_name, "'");

    input_tensors->emplace_back(node, idx);

    const auto& iter = input_nodes->find(node);
    if (iter == input_nodes->end()) {
      input_nodes->insert({node, {idx}});
    } else {
      auto& indices = iter->second;
      if (std::find(indices.begin(), indices.end(), idx) != indices.end()) {