auto sink = ops::Identity(scope.WithOpName("sink"), exit);

    // Remove the dummy node and add the loop backedge.
    scope.graph()->RemoveNode(dummy.node());
    scope.graph()->AddEdge(next_iteration.node(), 0, merge.output.node(), 1);

    TF_EXPECT_OK(scope.ToGraph(&graph));
  }

  GraphShapeInfo shape_info;
  TF_ASSERT_OK(InferShapes(&graph, /*arg_shapes=*/{}, /*fnlib_def=*/nullptr,

      << "Use of BlockArgument encounted before def!";
  // For argument, there is only one result output, so the index is always 0.
  graph_->AddEdge(input_node_it->second, 0, dst_node, dst_index);
  return absl::OkStatus();
}

Status Exporter::AddEdge(Operation* inst) {
  // For tf_executor.fetch, add only its data edges. Control edges are captured
  // later.

      << "Use of BlockArgument encounted before def!";
  // For argument, there is only one result output, so the index is always 0.
  graph_->AddEdge(input_node_it->second, 0, dst_node, dst_index);
  return absl::OkStatus();
}

Status Exporter::AddEdge(Operation* inst) {
  // For tf_executor.fetch, add only its data edges. Control edges are captured
  // later.

      TF_ASSIGN_OR_RETURN(placeholder_node, g->AddNode(placeholder_def));
      placeholders[placeholder_index] = placeholder_node;
    } else {
      placeholder_node = iter->second;
    }
    g->AddEdge(placeholder_node, 0, dst, dst_input);

    // Replace `e->dst()` because its input node changed.
    NodeDef new_def = dst->def();
    *new_def.mutable_input(dst_input) = placeholder_node->name();

    // TODO(sanjoy): This does not update NodeDef inputs.  To be able to update
    // NodeDef inputs we first need to fix encapsulate_subgraphs_pass to fix up
    // the NodeDef inputs to the function call nodes.
    g->AddEdge(new_node, edge->src_output(), edge->dst(), edge->dst_input());
    g->RemoveEdge(edge);
  }
}

// Returns a data value that is dead iff `control` is dead.
Output ControlToData(const Scope& scope, Node* control) {

      graph->AddEdge(data_inputs[i].first, data_inputs[i].second, xla_launch,
                     i);
    }
    for (Node* n : control_inputs) {
      graph->AddControlEdge(n, xla_launch);
    }
    for (int i = 0, end = data_outputs.size(); i < end; ++i) {
      for (const auto& successor : data_outputs[i]) {
        graph->AddEdge(xla_launch, i, successor.first, successor.second);
      }

  cloned_node->set_assigned_device_name(n->assigned_device_name());

  for (const Edge* in_edge : n->in_edges()) {
    graph->AddEdge(in_edge->src(), in_edge->src_output(), cloned_node,
                   in_edge->dst_input());
  }

  for (const Edge* out_edge_to_clone : out_edges_to_clone) {
    graph->AddEdge(cloned_node, out_edge_to_clone->src_output(),
                   out_edge_to_clone->dst(), out_edge_to_clone->dst_input());

      for (E edge : network.outEdges(node)) {
        N successorNode = network.incidentNodes(edge).adjacentNode(node);
        if (subgraph.nodes().contains(successorNode)) {
          subgraph.addEdge(node, successorNode, edge);
        }
      }
    }
    return subgraph;
  }

  /** Creates a mutable copy of {@code graph} with the same nodes and edges. */

  auto var =
      ops::VarHandleOp(root.WithOpName("var"), DT_INT32, TensorShape({}));
  auto int32_on_device =
      ops::ReadVariableOp(root.WithOpName("int32_on_device"), var, DT_INT32);

  root.graph()->AddEdge(int32_on_device.node(), 0, call, 0);

  std::unique_ptr<Graph> graph;
  TF_ASSERT_OK(BuildXlaOps(root, fdef_lib, &graph));

  Node* stateful_partitioned_call_op = nullptr;
  for (Node* n : graph->op_nodes()) {

      for (E edge : network.outEdges(node)) {
        N successorNode = network.incidentNodes(edge).adjacentNode(node);
        if (subgraph.nodes().contains(successorNode)) {
          subgraph.addEdge(node, successorNode, edge);
        }
      }
    }
    return subgraph;
  }

  /** Creates a mutable copy of {@code graph} with the same nodes and edges. */