mirrors/tinygrad
2
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-24 02:14:17 +00:00
Code Issues Projects Releases Packages Wiki Activity

fix dl/dr shape

Browse source
This commit is contained in:
George Hotz 2026-06-01 14:53:41 -07:00
commit bdbee57f34
1 changed files with 3 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

3
tinygrad/uop/ops.py
View file

@ -288,6 +288,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
# STAGE adds the existing shape to the front, opposite of INDEX
return tuple([int(r.vmax+1) for r in self.src[1:]])+self.src[0].shape
case Ops.DEFINE_LOCAL | Ops.DEFINE_REG:
if len(self.src) >= 1:
# NOTE: this is the same as PARAM
return tuple(self.src[0].sgep(i) for i in range(self.src[0].dtype.count))
if isinstance(self.dtype, PtrDType):
return (self.ptrdtype.size, self.dtype.count) if self.dtype.count > 1 else (self.ptrdtype.size,)
return (self.dtype.count,) if self.dtype.count > 1 else ()