import torch


def collect_engines(script_module, group_type):
    """
    Collect all engines in a script_module of the group type
    """
    def collect_engines_in_block(block):
        engine_names = [
            n.s('name') for n in block.nodes()
            if n.kind() == 'prim::GetAttr' and n.s('name').startswith(group_type)
        ]
        for n in block.nodes():
            for blk in n.blocks():
                engine_names += collect_engines_in_block(blk)
        return engine_names

    if isinstance(script_module, torch.jit.ScriptModule):
        script_module = script_module._c
    engine_names = collect_engines_in_block(script_module.forward.graph)
    engines = [script_module.getattr(name) for name in engine_names]
    return list(zip(engine_names, engines))