From a4a25720b2b25decda3b3126e67b313b3a318078 Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Thu, 15 May 2025 20:47:57 -0700
Subject: [PATCH] add test_multitensor_jit_input [pr] (#10347)

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

diff --git a/test/test_multitensor.py b/test/test_multitensor.py
index 0e16b6e17e..2455af8d7e 100644
--- a/test/test_multitensor.py
+++ b/test/test_multitensor.py
@@ -210,6 +210,14 @@ class TestMultiTensor(unittest.TestCase):
         a,b = jit_allreduce(Tensor.rand(256, 256))
         np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
 
+  def test_multitensor_jit_input(self):
+    @TinyJit
+    def f(x): return (x+1).contiguous().sum()
+    for _ in range(5):
+      tt = Tensor.arange(0, 4).contiguous().realize().shard((d1,d2), 0).realize()
+      out = f(tt)
+      assert out.item() == 1+2+3+4
+
   @unittest.skip("slow")
   def test_fuzz_allreduce(self):
     random.seed(41)