From 35e53c08095b84f6a3a2d616cc176aefb225cfd1 Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Tue, 11 Jun 2024 10:58:33 +0200
Subject: [PATCH] add sharded arange test (#4908)

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

diff --git a/test/test_multitensor.py b/test/test_multitensor.py
index 62ba69ba45..638de7dd72 100644
--- a/test/test_multitensor.py
+++ b/test/test_multitensor.py
@@ -48,6 +48,11 @@ class TestMultiTensor(unittest.TestCase):
       assert lb.shape == (128,)
     (X + X).realize()
 
+  def test_sharded_arange(self):
+    sharded_arange = Tensor.arange(1000).shard(devices_2, 0)
+    sharded_arange.realize()
+    np.testing.assert_equal(sharded_arange.numpy(), np.arange(1000))
+
   def test_shard_no_recompile(self):
     X = Tensor.ones(256).contiguous().realize()
     X.shard_((d0, d1), 0)