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

reduce_unparented only checks ranges (#12548)

Browse Source
This commit is contained in:
George Hotz
2025-10-09 14:14:03 +08:00
committed by GitHub
parent 375ee2c576
commit 80d99d52a5
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tinygrad/codegen/simplify.py
View File

@@ -129,7 +129,8 @@ def reduce_collapse(red:UOp):
def reduce_unparented(red:UOp):
if red.arg not in {Ops.ADD, Ops.MAX, Ops.MUL}: return None
reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].backward_slice_with_self)
assert all(x.op is Ops.RANGE for x in red.src[1:]), "some reduce srcs aren't ranges"
reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].ranges)
if len(reduce_unparented) == 0: return None
ret = red.replace(src=(red.src[0],)+tuple(reduce_parented)) if len(reduce_parented) or red.dtype != red.src[0].dtype else red.src[0]
if red.arg is Ops.ADD: