class_docstr: ``IterGraphModule`` provides the ability to do cross-iteration optimization.
  
  Given a ``fx.GraphModule``, main_gm, ``IterGraphModule`` internally
  duplicate it to 3 copies and redirect the ``forward`` request to a different
  ``fx.GraphModule`` based on the iteration count. This allows users to do
  graph optimizations that across iterations (e.g., moving collective wait in
  the backward to the forward of the next iteration).
  
  Note that users must call the APIs provided by ``IterGraphModule`` or
  ``IterGraph`` to rewrite the graph so that ``IterGraphModule`` can keep the
  data dependency for all 3 graphs.
class_docstr_tok: ["``","IterGraphModule","``","provides","the","ability","to","do","cross-iteration","optimization",".","Given","a","``","fx.GraphModule","``",",","main_gm",",","``","IterGraphModule","``","internally","duplicate","it","to","3","copies","and","redirect","the","``","forward","``","request","to","a","different","``","fx.GraphModule","``","based","on","the","iteration","count",".","This","allows","users","to","do","graph","optimizations","that","across","iterations","(","e.g.",",","moving","collective","wait","in","the","backward","to","the","forward","of","the","next","iteration",")",".","Note","that","users","must","call","the","APIs","provided","by","``","IterGraphModule","``","or","``","IterGraph","``","to","rewrite","the","graph","so","that","``","IterGraphModule","``","can","keep","the","data","dependency","for","all","3","graphs","."]
code: def print_readable(self, print_output: bool = True) -> str:
      return self.main_gm.print_readable(print_output)
  
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 8 for calling the IterGraphModule obj's underlying member method code with example usage: obj.print_readable(print_output) and returns: self