github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 07:28:15 -05:00
Code Issues Packages Projects Releases Wiki Activity

fixup grouper process replay [pr] (#10538)

Browse Source
This commit is contained in:
qazal
2025-05-27 20:46:55 +03:00
committed by GitHub
parent a07caaca0d
commit 617ecc1a7b
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tinygrad/engine/grouper.py
View File

@@ -580,7 +580,7 @@ def get_kernelize_map(big_sink:UOp) -> dict[UOp, UOp]:
with Context(PICKLE_BUFFERS=0):
import pickle
kernel_asts = [u.arg.ast for u in sched_sink.toposort() if u.op is Ops.KERNEL]
PROCESS_REPLAY_CAPTURE[id(big_sink)] = pickle.dumps((big_sink, ContextVar._cache, [u.arg.ast for u in kernel_asts if u.op is Ops.KERNEL]))
PROCESS_REPLAY_CAPTURE[id(big_sink)] = pickle.dumps((big_sink, ContextVar._cache, kernel_asts))
# map tensors to buffer/assign/const
# TODO: this is not right, and causes TestDataset.test_dataset_is_realized to fail unless I unprincipledly add Ops.COPY, which breaks others