mawei/tinygrad
1
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-13 00:15:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'master' into cstyle_new_style

Browse Source
This commit is contained in:
George Hotz
2026-06-06 09:49:46 -07:00
committed by GitHub
parent 1abbb51989 2a2f81dd3d
commit 489eb2cf72
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tinygrad/uop/ops.py
View File

@@ -696,15 +696,15 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
case Ops.STACK: return self.src[i].sintify()
case _: raise RuntimeError(f"no sgep on {self.op}")
@functools.cached_property
def as_shape(self) -> tuple[sint, ...]:
return tuple(ssimplify(self.sgep(i)) for i in range(max(self.dtype.count, len(self.src))))
@functools.cached_property
def marg(self):
match self.op:
case Ops.RESHAPE | Ops.EXPAND: return tuple(ssimplify(self.src[1].sgep(i)) for i in range(self.src[1].dtype.count))
case Ops.PAD | Ops.SHRINK:
# this is like broadcasting for shapes
return tuple(((ssimplify(self.src[1]) if self.src[1].shape == () else self.src[1].sgep(i)),
(ssimplify(self.src[2]) if self.src[2].shape == () else self.src[2].sgep(i)))
for i in range(max(self.src[1].dtype.count, self.src[2].dtype.count)))
case Ops.RESHAPE | Ops.EXPAND: return self.src[1].as_shape
case Ops.PAD | Ops.SHRINK: return tuple(zip(self.src[1].as_shape, self.src[2].as_shape))
case Ops.PERMUTE | Ops.FLIP: return self.arg
case _: raise RuntimeError(f"{self.op} is not a MovementOp")