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)