mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
12e95aa4ee
* Automatically applies ruff rule 401. Turns loops into equivalent list comprehensions which are faster and do not leak the scope of the loop variables. * list comprehensions not only often have better typing, but are 50+% faster than for loops on overhead. They also preserve length information etc and are better for the interpreter to optimize. * Manually went back and made mypy happy after the change. * Also fixed style lints in files covered by flake8 but not by pyfmt Pull Request resolved: https://github.com/pytorch/pytorch/pull/140980 Approved by: https://github.com/justinchuby, https://github.com/malfet
57 lines
1.3 KiB
Python
57 lines
1.3 KiB
Python
"""
|
|
Utilities for converting data types into structured JSON for dumping.
|
|
"""
|
|
|
|
import traceback
|
|
from typing import Any, Dict, List, Sequence, Set
|
|
|
|
import torch._logging._internal
|
|
|
|
|
|
INTERN_TABLE: Dict[str, int] = {}
|
|
|
|
|
|
DUMPED_FILES: Set[str] = set()
|
|
|
|
|
|
def intern_string(s: str) -> int:
|
|
r = INTERN_TABLE.get(s, None)
|
|
if r is None:
|
|
r = len(INTERN_TABLE)
|
|
INTERN_TABLE[s] = r
|
|
torch._logging._internal.trace_structured(
|
|
"str", lambda: (s, r), suppress_context=True
|
|
)
|
|
return r
|
|
|
|
|
|
def dump_file(filename: str) -> None:
|
|
if "eval_with_key" not in filename:
|
|
return
|
|
if filename in DUMPED_FILES:
|
|
return
|
|
DUMPED_FILES.add(filename)
|
|
from torch.fx.graph_module import _loader
|
|
|
|
torch._logging._internal.trace_structured(
|
|
"dump_file",
|
|
metadata_fn=lambda: {
|
|
"name": filename,
|
|
},
|
|
payload_fn=lambda: _loader.get_source(filename),
|
|
)
|
|
|
|
|
|
def from_traceback(tb: Sequence[traceback.FrameSummary]) -> List[Dict[str, Any]]:
|
|
# dict naming convention here coincides with
|
|
# python/combined_traceback.cpp
|
|
r = [
|
|
{
|
|
"line": frame.lineno,
|
|
"name": frame.name,
|
|
"filename": intern_string(frame.filename),
|
|
}
|
|
for frame in tb
|
|
]
|
|
return r
|