From f793cdeb871e7ce02bc5103f587f5d271c7f35f1 Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 9 Oct 2025 16:13:02 +0800 Subject: [PATCH] clean up shape changing logic to not use st [pr] (#12560) --- tinygrad/uop/ops.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index b7263bd9e2..99320939c6 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -216,16 +216,16 @@ class UOp(MathTrait, metaclass=UOpMetaClass): # otherwise we get the shape from sources if not (src_sts := [x.st for x in self.src if x.st is not None]): return None assert all_same([x.shape for x in src_sts]), f"UOp sources must have the same shape {self} {[x.shape for x in src_sts]}" + shape = src_sts[0].shape + # shape changing ops match self.op: - case Ops.MULTI: shape = tuple(s*len(self.device) if a == self.axis else s for a,s in enumerate(src_sts[0].shape)) + case Ops.MULTI: shape = tuple(s*len(self.device) if a == self.axis else s for a,s in enumerate(shape)) case Ops.BITCAST: - shape = src_sts[0].shape - if self.dtype.itemsize != (input_sz:=self.src[0].dtype.itemsize): shape = shape[:-1]+((shape[-1]*input_sz) // self.dtype.itemsize,) + if (output_sz:=self.dtype.itemsize) != (input_sz:=self.src[0].dtype.itemsize): shape = shape[:-1]+((shape[-1]*input_sz) // output_sz,) 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 = tuple(1 if i in axis_arg else s for i,s in enumerate(src_sts[0].shape)) - case _: shape = src_sts[0].shape + shape = tuple(1 if i in axis_arg else s for i,s in enumerate(shape)) return ShapeTracker.from_shape(shape) @property