Returns a list of submodules used for control flow operations that are in the given toplevel graph (does not look into submodules). Specifically, the returned value is a list containing tuples of (name of the submodule that's stored in the graph module, the submodule itself, and
(
graph_module: torch.fx.GraphModule,
op_to_submodule_arg_index: dict[HigherOrderOperator, list[int]],
)
| 49 | |
| 50 | |
| 51 | def _get_control_flow_submodules( |
| 52 | graph_module: torch.fx.GraphModule, |
| 53 | op_to_submodule_arg_index: dict[HigherOrderOperator, list[int]], |
| 54 | ) -> List[Tuple[str, torch.fx.GraphModule, torch.fx.Node]]: |
| 55 | """ |
| 56 | Returns a list of submodules used for control flow operations |
| 57 | that are in the given toplevel graph (does not look |
| 58 | into submodules). Specifically, the returned value is a list containing |
| 59 | tuples of (name of the submodule that's stored in the graph module, the |
| 60 | submodule itself, and the fx node that uses this submodule). |
| 61 | """ |
| 62 | control_flow_submodules = [] |
| 63 | for node in graph_module.graph.nodes: |
| 64 | if node.op != "call_function": |
| 65 | continue |
| 66 | |
| 67 | for op in op_to_submodule_arg_index: |
| 68 | if node.target is not op: |
| 69 | continue |
| 70 | for i in op_to_submodule_arg_index[op]: |
| 71 | control_flow_submodules.append(_get_submodule(graph_module, node, i)) |
| 72 | |
| 73 | return control_flow_submodules |
| 74 | |
| 75 | |
| 76 | def get_control_flow_submodules( |
no test coverage detected