class_docstr: A collection of remote modules (of type RemoteModule) with connections showing how inputs
  to the model or outputs of individual modules are use as inputs of subsequent modules.
  The graph has a number of helper functions that add new modules to the graph and define inputs
  to these module.
docstr: Splits the graph into pipeline partitions and for each parition returns a tuple (indices, module_rref),
  where indices is indices of modules of the partition in the graph, and module_rref is an RRef to an nn.Module:
  Each partition is a list of modules on the same device that are executed sequentially (output of each module is
  the input to the next module).
  If there is only one module in the partition, module_rref is reference to that module; otherwise those modules
  are wrapped by a MultiInputSequential and module_rref referes to that.
docstr_tok: ["Splits","the","graph","into","pipeline","partitions","and","for","each","parition","returns","a","tuple","(","indices",",","module_rref",")",",","where","indices","is","indices","of","modules","of","the","partition","in","the","graph",",","and","module_rref","is","an","RRef","to","an","nn.Module",":","Each","partition","is","a","list","of","modules","on","the","same","device","that","are","executed","sequentially","(","output","of","each","module","is","the","input","to","the","next","module",")",".","If","there","is","only","one","module","in","the","partition",",","module_rref","is","reference","to","that","module",";","otherwise","those","modules","are","wrapped","by","a","MultiInputSequential","and","module_rref","referes","to","that","."]
code: def partition_graph(self) -> List[Tuple[List[Node], rpc.RRef]]:
      """Splits the graph into pipeline partitions and for each parition returns a tuple (indices, module_rref),
      where indices is indices of modules of the partition in the graph, and module_rref is an RRef to an nn.Module:
      Each partition is a list of modules on the same device that are executed sequentially (output of each module is
      the input to the next module).
      If there is only one module in the partition, module_rref is reference to that module; otherwise those modules
      are wrapped by a MultiInputSequential and module_rref referes to that.
      """
      self._compile()
      modules_used: Set[Node] = set()
      partitions = []
      for node in self.nodes:
          if node in modules_used:
              continue
          partition = self._trace_modules(node)
          assert not modules_used.intersection(partition)
          modules_used.update(partition)
  
          if len(partition) == 1:
              remote_module = partition[0].module.get_module_rref()
          else:
              remote_module = rpc.remote(
                  partition[0].module.on,
                  RemoteSequential,
                  args=(
                      [p.module.get_module_rref() for p in partition],
                  ),
              )
          partitions.append((partition, remote_module))
  
      return partitions
  
imports: from dataclasses import dataclass
  from typing import List, Optional, Set, Tuple, Union
  from torch import Tensor, nn
  from torch.distributed import rpc
  from torch.distributed.nn import RemoteModule
  from .data import DataConsumer
label_desc: Use image node_id 9 for calling the PipelineModulesGraph obj's underlying member method code with example usage: obj.partition_graph() and returns: partitions