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","."]
code: def __init__(
      self,
      orig_graph: fx.Graph,
      setup_graph: fx.Graph,
      cleanup_graph: fx.Graph,
      owning_module: Optional[fx.GraphModule] = None,
      tracer_cls: Optional[Type["fx.Tracer"]] = None,
      tracer_extras: Optional[Dict[str, Any]] = None,
  ):
      super().__init__(owning_module, tracer_cls, tracer_extras)
  
      output_vals = self.graph_copy(
          orig_graph, {}, return_output_node=True
      )
      # TODO: if we do ``deepcopy(_codegen)`` and the input argument contains
      # a dictionary with the form of Dict[torch.Tensor, Any], the
      # torch.fx._pytree.treen_flatten_spec will not be able to flatten the
      # dict -- the torch.Tensor will be duplicated because the _input_spec
      # will save the ``keys`` of a dictionary (the values are not saved).
      self._codegen = copy.deepcopy(orig_graph._codegen)
      assert isinstance(output_vals, tuple)
      output_val, old_output_val = output_vals
      super().output(
          output_val, type_expr=getattr(old_output_val, "type", None)
      )
  
      self.setup_graph = setup_graph
      self.cleanup_graph = cleanup_graph
      self._all_graphs: Tuple[fx.Graph, ...] = (
          self.setup_graph,
          self.cleanup_graph,
          cast(fx.Graph, super()),
      )
  
      self._setup_mapping: Dict[fx.Node, fx.Node] = {}
      self._cleanup_mapping: Dict[fx.Node, fx.Node] = {}
      self._freeze_cross_iter_movement = False
      self._cross_iter_block_count = 0
  
      for node, setup_node, cleanup_node in zip(
          self.nodes, self.setup_graph.nodes, self.cleanup_graph.nodes
      ):
          self._setup_mapping[node] = setup_node
          self._cleanup_mapping[node] = cleanup_node
  
      self.num_extra_output = 0
  
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 1 to create a new IterGraph object from inherited base classes: fx with example: obj = IterGraph(orig_graph, setup_graph, cleanup_graph, owning_module, tracer_cls, tracer_extras)