From d1f9ade9a078a44a3ed1c2ede2e1b9c38e9df03b Mon Sep 17 00:00:00 2001 From: George Hotz Date: Wed, 29 Apr 2026 15:46:42 -0700 Subject: [PATCH] DEFINE_VAR can also have shape --- tinygrad/uop/ops.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 96f73fe742..4524142482 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -234,7 +234,6 @@ class UOp(OpMixin, metaclass=UOpMetaClass): if isinstance(self.src[0].dtype, PtrDType) and not isinstance(self.src[0].dtype, ImageDType) and not isinstance(self.dtype, PtrDType): return None - case Ops.STACK: return (len(self.src),) case Ops.INDEX: # if it has a vec dtype, set the shape if self.dtype.count > 1: return (self.dtype.count,) @@ -246,11 +245,11 @@ class UOp(OpMixin, metaclass=UOpMetaClass): return self.src[0].shape[len(self.src[1:]):] # some ops init the shape - case Ops.DEFINE_VAR | Ops.BIND | Ops.RANGE | Ops.SPECIAL: return () - case Ops.CONST: - # CONST can have shape if it has a vec dtype + case Ops.CONST | Ops.DEFINE_VAR: + # these can have shape if it has a vec dtype if self.dtype.count > 1: return (self.dtype.count,) return () + case Ops.BIND | Ops.RANGE | Ops.SPECIAL: return () case Ops.VCONST: return (len(self.arg),) case Ops.BUFFER: return (self.arg,) case Ops.BUFFER_VIEW: return (self.arg[0],)