fix symbolic sharded reshape (#17463)

test/unit/test_call.py

@@ -359,5 +359,15 @@ class TestCallMultiSharded(unittest.TestCase):
     np.testing.assert_allclose(a.grad.numpy(), b.numpy(), rtol=1e-5)
     np.testing.assert_allclose(b.grad.numpy(), a.numpy(), rtol=1e-5)
 
+  def test_symbolic_reshape_shard_axis(self):
+    toks = UOp.variable("toks", 1, 2).bind(2)
+    devs = ("CPU:0", "CPU:1")
+    x = Tensor(np.arange(16, dtype=np.float32).reshape(1, 2, 8)).shard(devs, axis=2).realize()
+    @function
+    def f(x:Tensor) -> Tensor: return x.reshape(1, x.shape[1], 2, 4)
+    out = f(x[:, :toks]).realize()
+    self.assertEqual(out.uop.axis, 2)
+    np.testing.assert_equal(out[:1, :2].to(devs[0]).numpy(), np.arange(16, dtype=np.float32).reshape(1, 2, 2, 4))
+
 if __name__ == '__main__':
   unittest.main()

tinygrad/schedule/multi.py

@@ -126,7 +126,7 @@ def reshape_multi(root:UOp, multi:UOp):
   new_shardings = []
   for ax, rng in multi.sharding:
     count = int(rng.vmax)+1
-    target = prod(multi.shape[:ax])
+    target = ssimplify(prod(multi.shape[:ax]))
     if target not in arg_acc: raise RuntimeError(f"reshape {multi.shape} -> {new_shape} moved items between shards")
     new_ax = len(arg_acc) - arg_acc[::-1].index(target) - 1
     if new_shape[new_ax] % count != 0: raise RuntimeError(f"reshape {multi.shape} -> {new_shape} moved items between shards")