From a9c47285dd26c20a8100cfe9eda16efcde30b8a5 Mon Sep 17 00:00:00 2001
From: David Hou <david.hou314@gmail.com>
Date: Wed, 21 Feb 2024 23:15:45 -0800
Subject: [PATCH] multitensor reshape tests

---
 test/test_multitensor.py | 47 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 47 insertions(+)

diff --git a/test/test_multitensor.py b/test/test_multitensor.py
index 92a22e7f40..3c01051f75 100644
--- a/test/test_multitensor.py
+++ b/test/test_multitensor.py
@@ -307,6 +307,53 @@ class TestMultiTensor(unittest.TestCase):
     # for i, ast in enumerate(asts):
     #   print(f"{i} {ast}")
 
+  def test_reshape_on_axis(self):
+    devices = (d0, d1, d2)
+
+    t0 = Tensor.rand((26, 15, 7)).shard(devices, axis=1)
+
+    # test split and rejoin to the right
+    t1 = t0.reshape((26, 3, 5, 7))
+    t2 = t0.reshape((26, 3, 35))
+    t3 = t1.reshape((26, 15, 7))
+    t4 = t2.reshape((26, 105,))
+
+    for t in [t0, t1, t2, t3, t4]:
+      assert t.lazydata.axis == 1
+      np.testing.assert_allclose(t.numpy().flatten(), t0.numpy().flatten())
+
+    # test shape-one axis
+    t5 = t4.reshape((26, 1, 105))
+    assert t5.lazydata.axis == 2
+
+    # test split and rejoin to the right and reshape to the left
+    t5 = t0.reshape((2, 13, 3, 5, 7))
+    t6 = t0.reshape((13, 2, 3, 7, 5))
+    t7 = t0.reshape((1, 13, 2, 3, 1, 7, 5))
+    np.testing.assert_allclose(t5.numpy().flatten(), t0.numpy().flatten())
+    assert t5.lazydata.axis == 2
+    np.testing.assert_allclose(t6.numpy().flatten(), t0.numpy().flatten())
+    assert t6.lazydata.axis == 2
+    np.testing.assert_allclose(t7.numpy().flatten(), t0.numpy().flatten())
+    assert t7.lazydata.axis == 3
+
+    # test no left join
+    with self.assertRaises((AssertionError, ValueError)):
+      t0.reshape((26*15,7))
+
+  def test_reshape_on_axis_uneven(self):
+    devices = (d0, d1, d2)
+    t0 = Tensor.rand((4, 8, 15)).shard(devices, axis=1)
+
+    # no split axis if uneven
+    with self.assertRaises((AssertionError, ValueError)):
+      t0.reshape((4,4,2,15))
+
+    # ok to split reshape left and right though
+    t1 = t0.reshape(2, 2, 8, 3, 5)
+    np.testing.assert_allclose(t0.numpy().flatten(), t1.numpy().flatten())
+    assert t1.lazydata.axis == 2
+
 @unittest.skipIf(CI and Device.DEFAULT in {"GPU", "CUDA", "METAL"}, "no GPU CI")
 class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
   # shrink a multitensor on sharded axis