parameters: (self, module_with_submodules: GraphModule)
return_statement: return dag
docstring: Return the dag structure and the new fx module with submodules.
docstring_summary: Return the dag structure and the new fx module with submodules.
docstring_tokens: Return the dag structure and the new fx module with submodules .
function: def dump_dag(self, module_with_submodules: GraphModule) -> DAG:
          """Return the dag structure and the new fx module with submodules."""
          dag = DAG()
          for node in module_with_submodules.graph.nodes:
              if node.op == "output":
                  break
              if node.op in {"placeholder", "get_attr"}:
                  continue
              if node.target == operator.__getitem__:
                  continue
              input_nodes: Dict[Node, None] = {}
              map_arg(node.args, lambda n: input_nodes.setdefault(n))
              map_arg(node.kwargs, lambda n: input_nodes.setdefault(n))
              # When a node has two or more output nodes,
              # it outputs its result to 'getitem' nodes.
              # Those 'getitem' nodes are the output node for this node.
              # Otherwise, the output node is this node itself.
              if len(node.users) > 1:
                  output_nodes = list(node.users)
              else:
                  output_nodes = [node]
              partition_id = int(node.name.rsplit("_", 1)[-1])
              device_ids = self.partitions[partition_id].logical_device_ids
              size_bytes = self.partitions[partition_id].used_mem_bytes
              dag.create_node(
                  node, list(input_nodes), output_nodes, device_ids, size_bytes
              )
          return dag
function_tokens: def dump_dag ( self , module_with_submodules : GraphModule ) -> DAG : dag = DAG ( ) for node in module_with_submodules . graph . nodes : if node . op == "output" : break if node . op in { "placeholder" , "get_attr" } : continue if node . target == operator . __getitem__ : continue input_nodes : Dict [ Node , None ] = { } map_arg ( node . args , lambda n : input_nodes . setdefault ( n ) ) map_arg ( node . kwargs , lambda n : input_nodes . setdefault ( n ) ) # When a node has two or more output nodes, # it outputs its result to 'getitem' nodes. # Those 'getitem' nodes are the output node for this node. # Otherwise, the output node is this node itself. if len ( node . users ) > 1 : output_nodes = list ( node . users ) else : output_nodes = [ node ] partition_id = int ( node . name . rsplit ( "_" , 1 ) [ - 1 ] ) device_ids = self . partitions [ partition_id ] . logical_device_ids size_bytes = self . partitions [ partition_id ] . used_mem_bytes dag . create_node ( node , list ( input_nodes ) , output_nodes , device_ids , size_bytes ) return dag