update TestSymbolicJit.test_plus1_pad (#12214)

was failing because movement was not captured
This commit is contained in:
chenyu
2025-09-16 15:57:50 -04:00
committed by GitHub
parent 2a72b00679
commit 6b808c5fe6

View File

@@ -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()