| |
| import sys |
| from typing import Optional |
|
|
| import torch |
| from torch._logging import LazyString |
|
|
|
|
| def lazy_format_graph_code(name, gm, maybe_id=None, **kwargs): |
| """ |
| Returns a LazyString that formats the graph code. |
| """ |
|
|
| def format_name(): |
| if maybe_id is not None: |
| return f"{name} {maybe_id}" |
| else: |
| return name |
|
|
| if "print_output" not in kwargs: |
| kwargs["print_output"] = False |
|
|
| if "colored" in kwargs: |
| try: |
| if not sys.stdout.isatty(): |
| kwargs["colored"] = False |
| except AttributeError: |
| kwargs["colored"] = False |
|
|
| return LazyString( |
| lambda: _format_graph_code( |
| f"===== {format_name()} =====\n", |
| gm.forward.__code__.co_filename, |
| gm.print_readable(**kwargs), |
| ) |
| ) |
|
|
|
|
| def _format_graph_code(name, filename, graph_str): |
| """ |
| Returns a string that formats the graph code. |
| """ |
| return f"TRACED GRAPH\n {name} {filename} {graph_str}\n" |
|
|
|
|
| def first_call_function_nn_module_stack(graph: torch.fx.Graph) -> Optional[dict]: |
| """ |
| Returns the nn_module_stack of the first call_function node. |
| """ |
| for node in graph.nodes: |
| if node.op == "call_function" and "nn_module_stack" in node.meta: |
| return node.meta["nn_module_stack"] |
| return None |
|
|
|
|
| def get_node_context(node, num_nodes=2) -> str: |
| """ |
| Returns a string of the last num_nodes nodes in the graph. |
| """ |
| node_contexts = [] |
| cur = node |
| for _ in range(num_nodes): |
| node_contexts.append(cur.format_node()) |
| if cur.op == "root": |
| break |
| cur = cur.prev |
| return "\n".join(node_contexts[::-1]) |
|
|
