class_docstr: ``IterGraph`` is used to perform cross-iteration optimization.
  
  ``IterGraph`` keeps track of the 3 graphs, self (the original graph), setup graph, and
  cleanup graph. The 3 graphs should be identical copies of a ``fx.Graph``.
  
  IterGraph subclass fx.Graph to override the necessary APIs that will be used
  when constructing a optimization, e.g., communication fusion. IterGraph also
  provides APIs that originally belong to fx.Node and all these APIs will have
  ``node_`` prefix. For example, ``IterGraph.node_prepend`` is the equivalence
  of ``fx.Node.prepend``. Note that all the optimizations must be constructed
  using these APIs.
class_docstr_tok: ["``","IterGraph","``","is","used","to","perform","cross-iteration","optimization",".","``","IterGraph","``","keeps","track","of","the","3","graphs",",","self","(","the","original","graph",")",",","setup","graph",",","and","cleanup","graph",".","The","3","graphs","should","be","identical","copies","of","a","``","fx.Graph","``",".","IterGraph","subclass","fx.Graph","to","override","the","necessary","APIs","that","will","be","used","when","constructing","a","optimization",",","e.g.",",","communication","fusion",".","IterGraph","also","provides","APIs","that","originally","belong","to","fx.Node","and","all","these","APIs","will","have","``","node_","``","prefix",".","For","example",",","``","IterGraph.node_prepend","``","is","the","equivalence","of","``","fx.Node.prepend","``",".","Note","that","all","the","optimizations","must","be","constructed","using","these","APIs","."]
docstr: Turn the inputs of a subgraph into the extra output of the entire graph.
  
  If ``erase_node`` is True, the subgraph will be erased from the graph -- essentially forward the inputs
  of the subgraph to the output of the graph.
code: def _forward_subgraph_inputs(
      self, subgraph: List[fx.Node], graph: fx.Graph, erase_node: bool
  ) -> int:
      """Turn the inputs of a subgraph into the extra output of the entire graph.
  
      If ``erase_node`` is True, the subgraph will be erased from the graph -- essentially forward the inputs
      of the subgraph to the output of the graph.
      """
      output = get_output(graph)
      inputs = []
      all_nodes: Set[fx.Node] = set(subgraph)
  
      for node in subgraph:
          node_inputs = pytree.arg_tree_leaves(
              *node.args, **node.kwargs
          )
          for _input in node_inputs:
              if not isinstance(_input, fx.Node):
                  continue
              if _input in all_nodes:
                  continue
              inputs.append(_input)
  
      if erase_node:
          # We have to remove the node in the reversed order to ensure the
          # node has zero users.
          erased = set()
          for node in reversed(subgraph):
              if len(node.users) == 1:
                  key = next(iter(node.users.keys()))
                  if key == output:
                      flatten_args, spec = tree_flatten(
                          (output.args, output.kwargs)
                      )
                      if node not in flatten_args:
                          # This optimizer node from the legacy _SPMD tracing.
                          node.users.clear()
                      elif str(node.target).startswith("aten.copy_"):
                          # This is the case where the optimizer is
                          # functionalized with copy_.
                          for i in range(len(flatten_args)):
                              if flatten_args[i] == node:
                                  flatten_args[i] = node.args[0]
                      else:
                          # We have not figured out semantics of forwarding
                          # all diff ops.
                          raise RuntimeError(
                              f"IterGraph does not how to forward the output of {node}"
                          )
                      output.args, output.kwargs = tree_unflatten(
                          flatten_args, spec
                      )
  
              # This is the step case where there is a virtual data dependency
              # (in-place update) between step and optimizer. And
              # functionalize_optim add this dependency
              for user in list(node.users.keys()):
                  if user in erased:
                      node.users.pop(user)
              if node.users:
                  raise RuntimeError(
                      "IterGraph has not supported moving the nodes that "
                      "produce users output result. "
                      f"Error node: {node}."
                  )
              self._fx_graph_call(graph, "erase_node", node)
              erased.add(node)
  
      # Add all the extra output nodes into a list and append the list to
      # the original output.args[0].
      if self.num_extra_output:
          # If the extra-output list already exist, just use it.
          cast(List[fx.Node], output.args[0][-1]).extend(inputs)  # type: ignore[index]
          new_output = output.args[0]
      else:
          # When adding the extra-output list, out_spec of _PyTreeCodeGen
          # must be updated accordingly.
          if isinstance(graph._codegen, _PyTreeCodeGen):
              codegen = graph._codegen
              new_output = list(output.args[0])  # type: ignore[arg-type]
              new_output.append(inputs)
              assert codegen.pytree_info.out_spec is not None
              original_tree_out = tree_unflatten(
                  cast(List[Any], output.args[0]),
                  codegen.pytree_info.out_spec,
              )
              # Use None as a placeholder. If we use the extra-output list
              # the list will be flatten as well and put into out_spec.
              _, out_spec = tree_flatten((original_tree_out, None))
              codegen.pytree_info = codegen.pytree_info._replace(
                  out_spec=out_spec
              )
          else:
              new_output = (output.args[0], inputs)
      self._fx_graph_call(graph, "erase_node", output)
      self._fx_graph_call(graph, "output", new_output)
  
      logger.info(
          "Extended outputs from the subgraph inputs: %s", str(inputs)
      )
      return len(inputs)
  
imports: import copy
  import inspect
  import logging
  from typing import Any, Callable, cast, Dict, List, Optional, Set, Tuple, Type
  import torch.nn
  from torch import fx
  from torch.distributed._spmd.graph_utils import clone_subgraph, get_output, is_leaf_subgraph
  from torch.distributed._spmd.partial_lower import partial_lower
  from torch.fx.graph import _PyTreeCodeGen, PythonCode
  from torch.fx.node import Argument
  from torch.profiler import record_function
  from torch.utils import _pytree
  from torch.utils._pytree import tree_flatten, tree_map, tree_map_only, tree_unflatten
label_desc: Use image node_id 7 for calling the IterGraph obj's underlying member method code with example usage: obj._forward_subgraph_inputs(subgraph, graph, erase_node) and returns: len