contiguous and same length
MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -3,7 +3,7 @@ from enum import Enum
from typing import Optional, Tuple, NamedTuple, Union, Any, List, Dict, Type
from copy import copy
import os, sys, functools, operator, weakref
-from tinygrad.helpers import ConvArgs, get_available_llops
+from tinygrad.helpers import ConvArgs, get_available_llops, prod
from tinygrad.shapetracker import ShapeTracker
# lazy can recurse a lot
@@ -251,7 +251,7 @@ class LazyBuffer:
# NOTE: if ret is in the cache, it can already be realized
if REMOVE_MOVEMENT_NOPS and ret.realized is None and x.realized is None and ret.st.contiguous:
root = get_lazybuffers(ret.op)[0]
- if root.st.contiguous and root != x:
+ if root.st.contiguous and root != x and prod(ret.st.shape) == prod(root.shape):
return root.movement_op(MovementOps.RESHAPE, ret.st.shape) if ret.st.shape != root.shape else root
return ret

MODIFY
tinygrad/shapetracker.py
tinygrad/shapetracker.py
@@ -102,7 +102,7 @@ class ShapeTracker:
def reshape(self, *new_shape):
assert all([isinstance(x, int) for x in new_shape])
- assert prod(self.shape) == prod(new_shape)
+ assert prod(self.shape) == prod(new_shape), f"can't reshape {self.shape} -> {new_shape}"
if self.shape == new_shape: return
# check if this is adding or removing 1s (only)

contiguous and same length