Function get_control_flow_submodules

exir/graph_module.py:76–93

Returns a list of submodules used for control flow operations (torch.ops.higher_order.cond/map/scan) 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 grap

(
    graph_module: torch.fx.GraphModule,
)

def get_control_flow_submodules(
    graph_module: torch.fx.GraphModule,
) -> List[Tuple[str, torch.fx.GraphModule, torch.fx.Node]]:
    """
    Returns a list of submodules used for control flow operations
    (torch.ops.higher_order.cond/map/scan) 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 the fx node that uses this submodule).
    """
    return _get_control_flow_submodules(
        graph_module,
        {
            torch.ops.higher_order.cond: [1, 2],
            torch.ops.higher_order.map_impl: [0],
            torch.ops.higher_order.scan: [0],  # combine_fn is at arg index 0
        },
    )


def get_cond_while_submodules(