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.
docstr: Move the ``subgraph`` to the next iteration before ``target_node``.
  
  The ``subgraph`` is a list of fx.Node and must satisfy the following
  restrictions:
      1. The order of the nodes in ``subgraph`` must obey the topological
         sort order.
      2. The users of the node in ``subgraph`` must be one of the following:
          a.) the user is also a node in ``subgraph``.
          b.) the user is the output of the full graph.
          c.) the node has users (side effect node).
docstr_tok: ["Move","the","``","subgraph","``","to","the","next","iteration","before","``","target_node","``",".","The","``","subgraph","``","is","a","list","of","fx.Node","and","must","satisfy","the","following","restrictions",":","1",".","The","order","of","the","nodes","in","``","subgraph","``","must","obey","the","topological","sort","order",".","2",".","The","users","of","the","node","in","``","subgraph","``","must","be","one","of","the","following",":","a",".",")","the","user","is","also","a","node","in","``","subgraph","``",".","b",".",")","the","user","is","the","output","of","the","full","graph",".","c.",")","the","node","has","users","(","side","effect","node",")","."]
code: def move_to_next_iter_before(
      self, subgraph: List[fx.Node], target_node: fx.Node
  ) -> None:
      """Move the ``subgraph`` to the next iteration before ``target_node``.
  
      The ``subgraph`` is a list of fx.Node and must satisfy the following
      restrictions:
          1. The order of the nodes in ``subgraph`` must obey the topological
             sort order.
          2. The users of the node in ``subgraph`` must be one of the following:
              a.) the user is also a node in ``subgraph``.
              b.) the user is the output of the full graph.
              c.) the node has users (side effect node).
      """
      if self._freeze_cross_iter_movement:
          raise RuntimeError(
              "The cross-iteration movement has been frozen for the given "
              "IterGraph."
          )
  
      if not is_leaf_subgraph(self, subgraph):
          raise ValueError(
              "The target nodes for ``move_to_next_iter_before`` must "
              "satisfy one of the following conditions: 1) the user of the "
              "node is in the target nodes, 2) the user is the ouput of the "
              "graph, 3) there are no users -- the node is a side-effect node. "
          )
  
      self._cross_iter_block_count += 1
      # The main graph must be the last one to be modified. Otherwise, the
      # mapping may change and hence introduce incorrect mapping for setup
      # and cleanup graphs.
  
      # For the setup graph, no additional input is needed but additional
      # outputs will be created. The additional output represents the input of
      # the action to be moved to the next iteration -- main graph.
      setup_subgraph: List[fx.Node] = []
      for node in subgraph:
          mapped_node = self._lookup_node(node, self.setup_graph)
          assert mapped_node is not None
          setup_subgraph.append(mapped_node)
      setup_extra_input = self._forward_subgraph_inputs(
          subgraph=setup_subgraph,
          graph=self.setup_graph,
          erase_node=True,
      )
  
      # For the cleanup graph, additional input is required to get the output
      # from the last iteration -- main graph. Additional nodes are also
      # needed to perform the action moved from the last iteration.
      target_cleanup_node = self._lookup_node(
          target_node, self.cleanup_graph
      )
      assert (
          target_cleanup_node is not None
      ), "The target_cleanup_node is None."
      cleanup_subgraph: List[fx.Node] = []
      for node in subgraph:
          mapped_node = self._lookup_node(node, self.cleanup_graph)
          assert mapped_node is not None
          cleanup_subgraph.append(mapped_node)
      cloned_subgraph = clone_subgraph(
          self.cleanup_graph,
          cleanup_subgraph,
          target=target_cleanup_node,
      )
      self._forward_inputs_to_subgraph(
          cloned_subgraph, self.cleanup_graph, setup_extra_input
      )
  
      # For the main graph, additional input will be created to represent
      # the output from the last iteration -- main graph or setup graph.
      # Additional output will also be generated to represent the input for
      # the next iteration -- the main graph or the cleanup graph.
      main_extra_input = self._forward_subgraph_inputs(
          subgraph=subgraph, graph=self, erase_node=False
      )
      assert main_extra_input == setup_extra_input
      for node in subgraph:
          target_node.prepend(node)
      self._forward_inputs_to_subgraph(subgraph, self, main_extra_input)
  
      # TODO: This is a temporary solution. We are going to remove DCE usage
      # or have something to replace fx DCE.
      for node in self.cleanup_graph.nodes:
          if len(node.users) == 0:
              node.users["__hold__"] = None  # type: ignore[index]
      for node in self.nodes:
          if len(node.users) == 0:
              node.users["__hold__"] = None  # type: ignore[index]
      self.num_extra_output += main_extra_input
  
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 9 for calling the IterGraph obj's underlying member method code with example usage: obj.move_to_next_iter_before(subgraph, target_node) without return types