diff --git a/test/test_symbolic_jit.py b/test/test_symbolic_jit.py index 1529e3bc21..881ce33489 100644 --- a/test/test_symbolic_jit.py +++ b/test/test_symbolic_jit.py @@ -197,6 +197,18 @@ 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_var_shape(self): + 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.ones(vi, 11).contiguous() + symbolic = a[:, 1:2] + symbolic = jf(symbolic).reshape(i, 1).numpy() + expected = f(a.reshape(i, 11)[:, 1: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) diff --git a/test/test_symbolic_ops.py b/test/test_symbolic_ops.py index 7bde68df45..ef74b45bc7 100644 --- a/test/test_symbolic_ops.py +++ b/test/test_symbolic_ops.py @@ -175,6 +175,14 @@ class TestSymbolicOps(unittest.TestCase): expected = a[3:5, i:i+2].numpy() np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) + def test_slice_var_shape(self): + for i in range(1, 5): + vi = Variable("i", 1, 10).bind(i) + a = Tensor.ones(vi, 11).contiguous() + symbolic = a[:, 1:2].reshape(i, 1).numpy() + expected = a.reshape(i, 11)[:, 1:2].numpy() + np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) + def test_ones_sum(self): for i in range(1, 5): vi = Variable("i", 1, 10).bind(i) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 544ddf8398..a575951f03 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1140,7 +1140,9 @@ class Tensor(MathTrait): boundary = [index, index+1] if index >= 0 else [index+size, index+size+1] case slice(): if index.step == 0: raise ValueError(f"{index=} cannot have 0 as step") - if all(isinstance(s, int) or s is None for s in (index.start,index.stop,index.step)): + if all(s is None for s in (index.start,index.stop,index.step)): + boundary, stride = [0, size], 1 + elif all(isinstance(s, int) or s is None for s in (index.start,index.stop,index.step)): # handle int slicing *boundary, stride = index.indices(cast(SupportsIndex, size)) if stride * (boundary[1] - boundary[0]) < 0: boundary = [0, 0]