forked from tinygrad/tinygrad
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:
@@ -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)
|
||||
|
||||
@@ -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}
|
||||
|
||||
Reference in New Issue
Block a user