xref: /aosp_15_r20/external/pytorch/torch/jit/generate_bytecode.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1# mypy: allow-untyped-defs
2from typing import List
3
4from torch._C import _compile_graph_to_code_table, _generate_upgraders_graph
5
6
7def format_bytecode(table):
8    # given a nested tuple, convert it to nested list
9    def listify(content):
10        if not isinstance(content, tuple):
11            return content
12        return [listify(i) for i in content]
13
14    formatted_table = {}
15    for entry in table:
16        identifier = entry[0]
17        content = entry[1]
18        content = listify(content)
19        formatted_table[identifier] = content
20    return formatted_table
21
22
23def generate_upgraders_bytecode() -> List:
24    yaml_content = []
25    upgraders_graph_map = _generate_upgraders_graph()
26    for upgrader_name, upgrader_graph in upgraders_graph_map.items():
27        bytecode_table = _compile_graph_to_code_table(upgrader_name, upgrader_graph)
28        entry = {upgrader_name: format_bytecode(bytecode_table)}
29        yaml_content.append(entry)
30    return yaml_content
31
32
33if __name__ == "__main__":
34    raise RuntimeError("This file is not meant to be run directly")
35