github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 15:08:02 -05:00
Code Issues Packages Projects Releases Wiki Activity

that had bugs, force an order (#2411)

Browse Source
This commit is contained in:
George Hotz
2023-11-23 15:52:16 -08:00
committed by GitHub
parent 65f4e6971b
commit 193be14b6c
5 changed files with 11 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
test/test_symbolic_jit.py
View File

@@ -55,7 +55,7 @@ class TestSymbolicJit(unittest.TestCase):
symbolic = jf(a.reshape(3, vi), b.reshape(vi, 5)).numpy()
expected = f(a, b).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert len(jf.jit_cache) == 2
assert len(jf.jit_cache) == 2 or getattr(Device[Device.DEFAULT], "graph", None)
def test_attention(self):
def f(q, k, v): return Tensor.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).realize()
@@ -68,7 +68,7 @@ class TestSymbolicJit(unittest.TestCase):
symbolic = jf(q, k.reshape(2, vi, 4, 8), v.reshape(2, vi, 4, 8)).reshape(2, 4, 1, 8).numpy()
expected = f(q, k, v).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert len(jf.jit_cache) == 6
assert len(jf.jit_cache) == 6 or getattr(Device[Device.DEFAULT], "graph", None)
def test_cat_dim0(self):
def f(a, b): return a.cat(b, dim=0).realize()