github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

_recurse_lb small cleanup (#4601)

Browse Source 
* minor cleanups

* comments

* extend env in replay
This commit is contained in:
qazal
2024-05-16 00:10:42 +08:00
committed by GitHub
parent 662bca8134
commit cd4d7e18c7
2 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
test/external/replay_schedule.py vendored
View File

@@ -1,5 +1,5 @@
#!/usr/bin/env python3
import subprocess, pickle, shlex, sys
import subprocess, pickle, shlex, sys, os
from typing import Dict, List, Tuple
from tinygrad.engine.graph import print_tree
from tinygrad.helpers import colored
@@ -12,7 +12,7 @@ def _run(name:str, cmd:List[str], env:Dict[str, str]) -> List[Tuple[LazyOp, ...]
return pickle.load(open(f"./{commit}.pkl", "rb"))
def _get_cmd():
parts, env = shlex.split(sys.argv[1]), {"SAVE_SCHEDULE": "1", "CAPTURE_AST": "1"}
parts, env = shlex.split(sys.argv[1]), {**os.environ, "SAVE_SCHEDULE": "1", "CAPTURE_AST": "1"}
env.update({k: v for p in parts if "=" in p for k, v in [p.split("=")]})
return [p for p in parts if "=" not in p], env