hotfix: test_jit_copyin

This commit is contained in:
George Hotz
2024-02-15 12:37:53 +01:00
parent 6356474d6d
commit ede4fd4705

View File

@@ -111,6 +111,15 @@ class TestJit(unittest.TestCase):
np.testing.assert_allclose(c.numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
assert_jit_cache_len(add_array, 1)
def test_jit_copyin(self):
@TinyJit
def f(a):
return a + Tensor([1,2,3])
for _ in range(5):
b = Tensor.randn(3)
c = f(b)
np.testing.assert_allclose(c.numpy(), b.numpy()+[1,2,3], atol=1e-4, rtol=1e-5)
def test_method_jit(self):
class Fun:
def __init__(self):