movement: resolve negative int slice bounds against symbolic sizes

negative int bounds in a slice against a symbolic dim were passed through
unresolved, giving wrong views. resolve them against the (possibly
symbolic) size, like slice.indices does for int dims
This commit is contained in:
2026-08-05 16:11:07 -07:00
parent 581bfdd94f
commit a8074f6e1b
2 changed files with 13 additions and 1 deletions
+9
View File
@@ -6,6 +6,15 @@ from examples.gpt2 import Attention
import numpy as np
class TestSymbolicOps(unittest.TestCase):
def test_negative_slice(self):
a = Tensor.rand(3, 10, 4)
for i in range(3, 10):
vi = Variable("i", 1, 10).bind(i)
# negative int bounds against a symbolic dim must resolve against the size, like slice.indices
np.testing.assert_allclose(a[:, :vi][:, -3:-1].numpy(), a[:, :i][:, -3:-1].numpy(), atol=1e-6, rtol=1e-6)
np.testing.assert_allclose(a[:, :vi][:, -1:].numpy(), a[:, :i][:, -1:].numpy(), atol=1e-6, rtol=1e-6)
np.testing.assert_allclose(a[:, :vi][:, -1].numpy(), a[:, :i][:, -1].numpy(), atol=1e-6, rtol=1e-6)
def test_plus1(self):
def f(a): return (a+1).realize()
a = Tensor.rand(3, 10)
+4 -1
View File
@@ -90,8 +90,11 @@ class MovementMixin:
if resolve(index.step == 0, False): raise ValueError(f"{index=} cannot have 0 as step")
start, stop = 0 if index.start is None else index.start, size if index.stop is None else index.stop
step = 1 if index.step is None else index.step
# resolve negative int bounds against the (possibly symbolic) size, like slice.indices
if isinstance(start, int) and start < 0: start = start + size
if isinstance(stop, int) and stop < 0: stop = stop + size
if all_int((start, stop, step)):
# handle int slicing (resolve negative bounds, clamp, stride)
# handle int slicing (clamp, stride)
*bound, stride = index.indices(int(size.vmax) if isinstance(size, UOp) else size)
bound = [0, 0] if stride * (bound[1] - bound[0]) < 0 else ([bound[1]+1, bound[0]+1] if stride < 0 else bound)
return {"size":ceildiv(bound[1]-bound[0], abs(stride)), "boundary":tuple(bound), "stride":stride, "collapse_dim":False}