diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py index cf77736e19626..0f05188013f64 100644 --- a/tinygrad/lazy.py +++ b/tinygrad/lazy.py @@ -277,7 +277,7 @@ def splitted_shape(dim_aft_div): return self.shape[:dim_to_split] + (self.shape[ # *** movement ops *** def _movement_op(self, st: ShapeTracker, op: MovementOps, arg: Union[Tuple[sint, ...], Tuple[Tuple[sint, sint], ...]]) -> LazyBuffer: - if SHUFFLE_MOVEMENT_OPS and self.optype == BinaryOps and not self.realized and (op in {MovementOps.SHRINK, MovementOps.STRIDE, MovementOps.PERMUTE} or (op == MovementOps.RESHAPE and self.op.op in UnaryOps)) and not self.children: + if SHUFFLE_MOVEMENT_OPS and self.optype == BinaryOps and not self.realized and (op in {MovementOps.SHRINK, MovementOps.STRIDE, MovementOps.PERMUTE} or (op == MovementOps.RESHAPE)) and not self.children: # and self.op.op in UnaryOps return self.op.replace_with_movement_ops([(op, arg)]) if REMOVE_MOVEMENT_NOPS and not self.realized and st.contiguous: # MovementOps aren't stacked any more, they each have one parent, find the root