fix issue with multi flip (#13115)

2025-11-05 15:28:50 -08:00
parent 4027eef264
commit 9b2b535fa4
3 changed files with 8 additions and 1 deletions
@@ -596,6 +596,12 @@ class TestMultiTensor(unittest.TestCase):
    # ast are the same on devices
    self.assertEqual(len(set(asts)), 1)

+  def test_flip(self):
+    rng = Tensor.rand((10, 10, 10))
+    t0 = rng.shard(devices_2, axis=1)
+    out = t0.flip(0) + 1
+    self.assertTrue((rng.flip(0)+1).allclose(out.to(rng.device)))
+
  def test_reshape_on_axis(self):
    t0 = Tensor.rand((26, 15, 7)).shard(devices_3, axis=1)