Function get_scan_submodules

exir/graph_module.py:115–133

Returns a list of submodules used for scan operations (torch.ops.higher_order.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

(
    graph_module: torch.fx.GraphModule,
)

def get_scan_submodules(
    graph_module: torch.fx.GraphModule,
) -> List[Tuple[str, torch.fx.GraphModule, torch.fx.Node]]:
    """
    Returns a list of submodules used for scan operations
    (torch.ops.higher_order.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).

    For scan, the combine_fn submodule is at argument index 0.
    The scan operator signature is: scan(combine_fn, init, xs, additional_inputs)
    """
    return _get_control_flow_submodules(
        graph_module,
        {
            torch.ops.higher_order.scan: [0],
        },
    )


def bfs_trace_with_node_process(