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

simple symbolic slice in llama [pr] (#10112)

Browse Source 
support slice that has step None and stop > start
This commit is contained in:
chenyu
2025-04-30 14:36:35 -04:00
committed by GitHub
parent b583ece8f3
commit 17d4d258ea
4 changed files with 38 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
test/test_symbolic_jit.py
View File

@@ -184,6 +184,19 @@ class TestSymbolicJit(unittest.TestCase):
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_slice(self):
# slice is a movement, so we pair it with a simple function to test the JIT interaction
def f(a): return (a+1).realize()
jf = TinyJit(f)
for i in range(1, 5):
vi = Variable("i", 1, 10).bind(i)
a = Tensor.rand(7, 11)
symbolic = a[3:5, vi:vi+2]
symbolic = jf(symbolic).numpy()
expected = f(a[3:5, i:i+2]).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
assert_jit_cache_len(jf, 1)
def test_ones_sum(self):
def f(a): return a.sum().realize()
jf = TinyJit(f)