diff --git a/tinygrad/shape/shapetracker.py b/tinygrad/shape/shapetracker.py
index 2f04cea468..9435b909f9 100644
--- a/tinygrad/shape/shapetracker.py
+++ b/tinygrad/shape/shapetracker.py
@@ -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)
 
diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py
index 1ed35ab229..b7263bd9e2 100644
--- a/tinygrad/uop/ops.py
+++ b/tinygrad/uop/ops.py
@@ -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)