minor reshape cleanups (#2619)

* minor reshape cleanups

* mea culpa

test/unit/test_shapetracker.py

@@ -340,6 +340,7 @@ class TestIndexExpressions2d(unittest.TestCase):
     v = self.st.views[-1]
     assert v.strides == (0, 0, 0, 1, 0, 0) and v.mask == ((0, 1), (0, 1), (0, 1), (30, 75), (0, 1), (0, 1)) and v.offset == -30
     self.st.assert_same()
+
 class TestSimplifyingShapeTracker(unittest.TestCase):
   def setUp(self):
     self.st = CheckingShapeTracker((1, 10))

tinygrad/shape/shapetracker.py

@@ -4,7 +4,6 @@ import functools, itertools, operator
 from dataclasses import dataclass
 from typing import Tuple, List, Optional, Dict, Set, cast, Union, Iterable
 from tinygrad.ops import MovementOps
-
 from tinygrad.helpers import prod, DEBUG, merge_dicts
 from tinygrad.shape.symbolic import Variable, MulNode, Node, SumNode, NumNode, sint
 from tinygrad.shape.view import View, _merge_dims

tinygrad/shape/view.py

@@ -28,44 +28,45 @@ def _merge_dims(shape:Tuple[int, ...], strides:Tuple[int, ...], mask:Optional[Tu
       ret[-1] = (ret[-1][0] * sh, st, (sh if state == 1 else ret[-1][2] * sh) if st else 0)
     else: ret.append((sh, st, sh if st else 0)) # begin new
     # merging ends with either non-zero strided dim or zero strided dim with mask range > 1
-    state = 1 if mask and st == 0 and mask[i][1] - mask[i][0] == 1 else (2 if state != 0 else 0)
+    state = 1 if (st == 0 and mask and mask[i][1] - mask[i][0] == 1) else (2 if state != 0 else 0)
   return tuple(ret)
 
 @functools.lru_cache(maxsize=None)
 def _reshape_mask(view: View, new_shape:Tuple[sint, ...]) -> Tuple[Optional[Tuple[Tuple[sint, sint], ...]], Optional[Tuple[sint, ...]], bool]:
-  if view.mask is None: return view.mask, tuple(), False
+  if view.mask is None: return view.mask, None, False
   new_mask: List[Tuple[int, int]] = []
 
   r_masks, r_shape, r_new_shape = reversed(view.mask), reversed(view.shape), reversed(new_shape)
   curr_stride, off, offsets, old_dim, new_dim, mask = 1, 0, [], next(r_shape, 1), next(r_new_shape, 1), next(r_masks, (0,1))
   #  off represents offset while combining masks of range one & zero stride
-  if mask[1] - mask[0] < 1: return ((0, 0),) * len(new_shape), tuple(), False # invalid mask
+  if mask[1] - mask[0] < 1: return ((0, 0),) * len(new_shape), None, False # invalid mask
 
   while len(new_mask) < len(new_shape):
-    (l, r), next_stride = (mask[0], mask[1]), new_dim * curr_stride
+    (l, r), next_stride = mask, new_dim * curr_stride
 
-    if old_dim >= new_dim * curr_stride: # need to split mask.
+    if old_dim >= next_stride: # need to split mask.
       offsets.append(off)
 
       if old_dim == next_stride: # simply copy the mask and get next batch for merging
         new_mask.append((l // curr_stride, (r - 1) // curr_stride + 1))
         curr_stride, off, old_dim, new_dim, mask = 1, 0, next(r_shape, 1), next(r_new_shape, 1), next(r_masks, (0,1))
-        if mask[1] - mask[0] < 1: return ((0, 0),) * len(new_shape), tuple(), False # invalid mask
+        if mask[1] - mask[0] < 1: return ((0, 0),) * len(new_shape), None, False # invalid mask
 
       else: # mask can only be splitted if reshape doesn't cut across the mask.
-        if ((l % (ns := next_stride) != 0 or r % ns != 0) and l // ns != (r - 1) // ns): return view.mask, tuple(), True
-        new_mask.append((l % ns // curr_stride, (r - 1) % ns // curr_stride + 1))
+        if ((l % next_stride != 0 or r % next_stride != 0) and l // next_stride != (r - 1) // next_stride): return view.mask, None, True
+        new_mask.append((l % next_stride // curr_stride, (r - 1) % next_stride // curr_stride + 1))
         curr_stride, new_dim = next_stride,  next(r_new_shape, 1) # need to get mask for next dimension
 
     else:
       next_mask = next(r_masks, (0, 1))
       # combine if the mask can unfold continuously
-      if (l != 0 or r != old_dim) and next_mask[1] - next_mask[0] != 1: return view.mask, tuple(), True
+      if mask != (0, old_dim) and next_mask[1] - next_mask[0] != 1: return view.mask, None, True
       if next_mask != (0, 1) and mask != (0, 1) and (next_mask[1] - next_mask[0] == 1): off += next_mask[0] * old_dim
       mask, old_dim = (next_mask[0] * old_dim + l, (next_mask[1] - 1) * old_dim + r), old_dim * next(r_shape, 1)
 
+  # TODO: can this be removed?
   for mask in r_masks: # if the old shape has leading 1s, need to make sure their mask is (0,1)
-    if mask != (0, 1): return ((0, 0),) * len(new_shape), tuple(), False
+    if mask != (0, 1): return ((0, 0),) * len(new_shape), None, False
 
   return tuple(reversed(new_mask)), tuple(offsets), False
 
@@ -171,7 +172,7 @@ class View:
     for merged_dim, s, real_dim in reversed(_merge_dims(self.shape, self.strides, self.mask)):
       acc, new_stride = 1, s
       while acc <= merged_dim and acc != merged_dim and (new_dim := next(r_new_shape, None)):
-        strides.append(new_stride if new_dim !=1 else 0)
+        strides.append(new_stride if new_dim != 1 else 0)
         if new_dim == 1: continue
         new_stride *= (new_dim if (acc :=  acc * new_dim) < real_dim else 0)
       if acc != merged_dim: break