move some ifs from merge_dims to reshape [pr] (tinygrad#8229)

the third return value is only used in reshape

test/unit/test_view.py

@@ -50,7 +50,7 @@ def test_0_in_strides(self):
     shape = (2, 3, 4)
     self.assertEqual(merge_dims(shape, (0, 4, 1)), ((2, 0, 0), (12, 1, 12)))
     self.assertEqual(merge_dims(shape, (0, 0, 1)), ((6, 0, 0), (4, 1, 4)))
-    self.assertEqual(merge_dims(shape, (3, 1, 0)), ((6, 1, 6), (4, 0, 0)))
+    self.assertEqual(merge_dims(shape, (3, 1, 0)), ((6, 1, 6), (4, 0, 4)))
     self.assertEqual(merge_dims(shape, (0, 0, 0)), ((24, 0, 0),))
 
   def test_pad_reshape(self):
@@ -59,16 +59,28 @@ def test_pad_reshape(self):
     # shift mask on stride 0
     self.assertEqual(merge_dims((2, 3), (0, 1), ((0, 1), (0, 2))), ((6, 1, 3),))
     # permute 0 / 1
-    self.assertEqual(merge_dims((3, 2), (1, 0), ((0, 2), (1, 2))), ((3, 1, 3), (2, 0, 0)))
+    self.assertEqual(merge_dims((3, 2), (1, 0), ((0, 2), (1, 2))), ((3, 1, 3), (2, 0, 2)))
+
+    # st = ShapeTracker.from_shape((1, 1, 2)).pad(((1, 0), (1, 0), (0, 1)))
+    # print(f"{st.views[-1]}")
+    self.assertEqual(merge_dims((2, 2, 3), (0, 0, 1), ((1, 2), (1, 2), (0, 2))), ((12, 1, 3),))
+
+    # st = ShapeTracker.from_shape((1, 1, 2, 2)).pad(((1, 0), (1, 0), (0, 1), (0, 1)))
+    # print(f"{st.views[-1]}")
+    self.assertEqual(merge_dims((2, 2, 3, 3), (0, 0, 2, 1), ((1, 2), (1, 2), (0, 2), (0, 2))), ((12, 2, 3), (3, 1, 3)))
+
+    # st = ShapeTracker.from_shape((1, 1, 1, 1)).pad(((0, 2), (0, 0), (0, 1), (0, 3)))
+    # print(f"{st.views[-1]}")
+    # self.assertEqual(merge_dims((3, 1, 2, 4), (0, 0, 0, 0), ((0, 1), (0, 1), (0, 1), (0, 1))), ((24, 0, 0),))
 
   def test_different_1_pad(self):
     # st = ShapeTracker.from_shape((2, 2, 1)).pad(((0, 0), (0, 0), (0, 1)))
     # print(f"{st.views[-1]}")
-    self.assertEqual(merge_dims((2, 2, 2), (2, 1, 0), ((0, 2), (0, 2), (0, 1))), ((4, 1, 4), (2, 0, 0)))
+    self.assertEqual(merge_dims((2, 2, 2), (2, 1, 0), ((0, 2), (0, 2), (0, 1))), ((4, 1, 4), (2, 0, 2)))
 
     # st = ShapeTracker.from_shape((2, 1, 1)).pad(((0, 0), (0, 1), (0, 1)))
     # print(f"{st.views[-1]}")
-    self.assertEqual(merge_dims((2, 2, 2), (1, 0, 0), ((0, 2), (0, 2), (0, 1))), ((2, 1, 2), (4, 0, 0)))
+    self.assertEqual(merge_dims((2, 2, 2), (1, 0, 0), ((0, 2), (0, 2), (0, 1))), ((2, 1, 2), (4, 0, 4)))
 
 if __name__ == '__main__':
   unittest.main()

tinygrad/shape/view.py

@@ -29,8 +29,8 @@ def merge_dims(shape:Tuple[int, ...], strides:Tuple[int, ...], mask:Optional[Tup
     if s == 1: continue
     last_s, last_st, last_pre_expand_s = ret[-1]
     # merge last dim with this dim if merging or strides matched
-    if merging or last_st == s * st: ret[-1] = (last_s * s, st, (s if merging else last_pre_expand_s * s) if st != 0 else 0)
-    else: ret.append((s, st, s if st != 0 else 0))
+    if merging or last_st == s * st: ret[-1] = (last_s * s, st, (s if merging else last_pre_expand_s * s))
+    else: ret.append((s, st, s))
     # merge this dim to next dim if size is 1
     merging = (mask[i][1] - mask[i][0] == 1) if mask is not None else s == 1
   return tuple(ret)
@@ -321,13 +321,13 @@ def reshape(self, new_shape: Tuple[sint, ...]) -> Optional[View]:
 
     r_strides, r_new_shape = [], reversed(new_shape)
     for merged_size, new_stride, real_size in reversed(merge_dims(self.shape, self.strides, self.mask)):
+      if new_stride == 0: real_size = 0
       # TODO: write with get_contraction
       acc = 1
       # TODO: third resolve shouldn't be needed
       while resolve(acc <= merged_size) and resolve(acc != merged_size) and resolve((new_dim := next(r_new_shape, 0)) > 0):
         r_strides.append(new_stride * acc)
         acc = acc * new_dim
-        # merge dim merges if (1) previous_stride = stride * dim, (2) dim = 1, stride = 0, either padded or not
         if not resolve(acc < real_size): new_stride = 0
       if resolve(acc != merged_size): return None