mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-02-07 05:05:13 -05:00
29 lines
1.3 KiB
Python
29 lines
1.3 KiB
Python
from dataclasses import dataclass
|
|
import difflib, logging, traceback, subprocess
|
|
from typing import Dict, Optional
|
|
from tinygrad.helpers import ContextVar, colored, getenv
|
|
|
|
def print_diff(s0, s1, unified=getenv("UNIFIED_DIFF",1)):
|
|
if not logging.getLogger().hasHandlers(): logging.basicConfig(level=logging.INFO, format="%(message)s")
|
|
if unified:
|
|
lines = list(difflib.unified_diff(str(s0).splitlines(), str(s1).splitlines()))
|
|
diff = "\n".join(colored(line, "red" if line.startswith("-") else "green" if line.startswith("+") else None) for line in lines)
|
|
else:
|
|
import ocdiff
|
|
diff = ocdiff.console_diff(str(s0), str(s1))
|
|
logging.info(diff)
|
|
|
|
@dataclass(frozen=True)
|
|
class ProcessReplayContext:
|
|
ctx_vars: Dict[str, int]
|
|
loc: str = ""
|
|
head_sha: str = ""
|
|
run_id: Optional[int] = None
|
|
def get_process_replay_ctx() -> ProcessReplayContext:
|
|
stack = filter(lambda x: "tinygrad" in x.filename and not any(n in x.filename for n in ["engine/schedule.py", "engine/realize.py", \
|
|
"codegen/kernel.py", "unittest"]), traceback.extract_stack()[:-1])
|
|
loc = "\n".join(traceback.format_list(stack))
|
|
try: head_sha = subprocess.check_output(["git", "rev-parse", "HEAD"]).strip().decode()
|
|
except Exception: head_sha = ""
|
|
return ProcessReplayContext({k:v.value for k,v in ContextVar._cache.items()}, loc, head_sha, getenv("GITHUB_RUN_ID") or None)
|