diff --git a/test/test_symbolic_jit.py b/test/test_symbolic_jit.py index f8dfbdfc31..f312539e7a 100644 --- a/test/test_symbolic_jit.py +++ b/test/test_symbolic_jit.py @@ -16,9 +16,9 @@ class TestSymbolicJit(unittest.TestCase): np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) assert_jit_cache_len(jf, 1) - @unittest.expectedFailure # TODO: fix, this works without jit def test_plus1_pad(self): - def f(a): return (a+1).pad((None, (0, 10-a.shape[1]))).realize() + # TODO: without contiguous, the pad is not captured in jit + def f(a): return (a+1).pad((None, (0, 10-a.shape[1]))).contiguous().realize() jf = TinyJit(f) a = Tensor.rand(3, 10) for i in range(1, 5): @@ -26,7 +26,7 @@ class TestSymbolicJit(unittest.TestCase): symbolic = jf(a[:, :vi]).numpy() expected = f(a[:, :i]).numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) - assert_jit_cache_len(jf, 1) + assert_jit_cache_len(jf, 2) # one add and one pad, can be one kernel? def test_add(self): def f(a, b): return (a+b).realize()