From cddfd8e25da3fdcf4211053f7dbe2138d514299c Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Fri, 12 Jul 2024 17:02:40 -0700 Subject: [PATCH] bugfix: group for reduce should check all dimensions (#5431) --- tinygrad/codegen/lowerer.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tinygrad/codegen/lowerer.py b/tinygrad/codegen/lowerer.py index 65398ea9c6..a4be11e0e2 100644 --- a/tinygrad/codegen/lowerer.py +++ b/tinygrad/codegen/lowerer.py @@ -96,9 +96,7 @@ class Lowerer(Kernel): if x.arg.idx != -1: has_valid = True for oidx, ridx in zip(self.idxs, self.ridxs): - if oidx != ridx: - valid = valid * oidx.eq(0) - break # TODO: this is wrong, but should maintain process replay + if oidx != ridx: valid = valid * oidx.eq(0) return UOp(UOps.STORE, None, (buf, idx, self.to_uop(x.src[0])) + ((valid,) if has_valid else ())) in_uops = tuple(self.to_uop(y) for y in x.src)