| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667 |
- # mypy: allow-untyped-defs
- 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])
|
