Returns a list of submodules used for control flow operations (torch.ops.higher_order.cond/while_loop) 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 gr
(
graph_module: torch.fx.GraphModule,
)
| 94 | |
| 95 | |
| 96 | def get_cond_while_submodules( |
| 97 | graph_module: torch.fx.GraphModule, |
| 98 | ) -> List[Tuple[str, torch.fx.GraphModule, torch.fx.Node]]: |
| 99 | """ |
| 100 | Returns a list of submodules used for control flow operations |
| 101 | (torch.ops.higher_order.cond/while_loop) that are in the given toplevel graph (does not look |
| 102 | into submodules). Specifically, the returned value is a list containing |
| 103 | tuples of (name of the submodule that's stored in the graph module, the |
| 104 | submodule itself, and the fx node that uses this submodule). |
| 105 | """ |
| 106 | return _get_control_flow_submodules( |
| 107 | graph_module, |
| 108 | { |
| 109 | torch.ops.higher_order.cond: [1, 2], |
| 110 | torch.ops.higher_order.while_loop: [0, 1], |
| 111 | }, |
| 112 | ) |
| 113 | |
| 114 | |
| 115 | def get_scan_submodules( |
nothing calls this directly
no test coverage detected