From cfdff84df07c3c6d026207f1d683775fa9fdf42b Mon Sep 17 00:00:00 2001 From: George Hotz Date: Wed, 29 Apr 2026 18:20:20 -0700 Subject: [PATCH] fixes --- tinygrad/schedule/rangeify.py | 2 +- tinygrad/uop/ops.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index f1978c8423..5be5b44dc6 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -546,7 +546,7 @@ def split_store(x:UOp) -> UOp|None: # if we have any open ranges here, we don't split if x.ranges: return None # raw STORE (not from bufferize_to_store) should be processed through its END wrapper, not independently - if x.op is Ops.STORE and x.src[0]._shape is not None: return None + #if x.op is Ops.STORE and x.src[0]._shape is not None: return None # local kernel rewrite lctx = LocalAddBufferContext() diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 758fb21d4c..546abf1e68 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -100,6 +100,8 @@ class UOpMetaClass(type): buffers[created] = _buffer if SPEC > 1: from tinygrad.uop.spec import full_spec, test_pyrender + # NOTE: while indexing, the STORE is wrong when only one side is updated + #shp = created._shape if SPEC > 2: test_pyrender(created) with Context(CHECK_OOB=0): fret = cast(bool|None, full_spec.rewrite(created)) if fret is not True: raise RuntimeError(f"SPEC ISSUE {fret}: {created}") @@ -237,7 +239,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass): case Ops.GEP: return (len(self.arg),) if len(self.arg) > 1 else () case Ops.STACK: return (len(self.src),) case Ops.INDEX: - shp = [] + shp:list[sint] = [] + # NOTE: the acc buffer can have a dtype with count, we need it back here + if self.src[0].dtype.count > 1: shp.append(self.src[0].dtype.count) for s in self.src[1:]: shp.extend(list(s.shape)) return tuple(shp) + self.src[0].shape[len(self.src[1:]):]