sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

remove ShapeTracker.reduce [pr] (#12559)

Browse Source
This commit is contained in:
chenyu
2025-10-09 03:54:11 -04:00
committed by GitHub
parent c1cc277fc3
commit 1bcea19846
2 changed files with 1 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/shape/shapetracker.py
-2
View File
@@ -53,8 +53,6 @@ class ShapeTracker:
@property
def size(self) -> int: return self.views[-1].size()
def reduce(self, axis:tuple[int, ...]) -> tuple[sint, ...]: return tuple(1 if i in axis else s for i,s in enumerate(self.shape))
def to_valid_uop(self, _idxs:list[UOp]|tuple[UOp, ...]|None=None) -> UOp:
return views_to_valid_uop(self.views, tuple(_idxs) if _idxs is not None else None)
tinygrad/uop/ops.py
+1 -1
View File
@@ -224,7 +224,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
case Ops.REDUCE_AXIS | Ops.WMMA:
axis_arg = self.arg[1] if self.op is Ops.REDUCE_AXIS else self.arg[7]
assert isinstance(axis_arg, tuple) and all(isinstance(x, int) for x in axis_arg), f"invalid type for axis: {axis_arg}"
shape = src_sts[0].reduce(axis_arg)
shape = tuple(1 if i in axis_arg else s for i,s in enumerate(src_sts[0].shape))
case _: shape = src_sts[0].shape
return ShapeTracker.from_shape(shape)