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

remove disk special case in mstack_early_shrink [pr] (#14494)

Browse Source
This commit is contained in:
chenyu
2026-02-02 08:34:45 -05:00
committed by GitHub
parent 965940dd00
commit a908f447d5
1 changed files with 1 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tinygrad/schedule/multi.py
View File

@@ -94,12 +94,7 @@ def mstack_early_shrink(ms:UOp, shrink:UOp):
return s.shrink(tuple(new_arg))
for i, x in enumerate(ms.src):
if x.op is Ops.COPY:
# if src device doesn't have a renderer, we have to view after the copy
# TODO: a way to understand this
if x.src[0].device in {"DISK", "NPY"}:
ret.append(apply_shrink(x, i))
else:
ret.append(apply_shrink(x.src[0], i).copy_to_device(x.device))
ret.append(apply_shrink(x.src[0], i).copy_to_device(x.device))
else:
ret.append(apply_shrink(x, i).contiguous())
return ms.replace(src=tuple(ret))