Allow negative step in index (#928)

Author: tomwhite

.github/workflows/array-api-tests.yml

@@ -122,8 +122,6 @@ jobs:
           array_api_tests/test_signatures.py::test_array_method_signature[__setitem__]
           array_api_tests/test_signatures.py::test_array_method_signature[to_device]
 
-          # (getitem with negative step size is not implemented)
-          array_api_tests/test_array_object.py::test_getitem
           # (indexing with more than one integer index is not implemented)
           array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_1
           array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_2

cubed/core/indexing.py

@@ -43,11 +43,19 @@ def index(x, key):
         if n:
             where_newaxis[i] -= n
     idx = ndindex.Tuple(*(ia for ia in idx.args if not isinstance(ia, ndindex.Newaxis)))
-    selection = idx.raw
+    selection = list(idx.raw)
+
+    # Use trick from xarray for negative step values
+    where_negative_step = []
+    for i, ia in enumerate(idx.args):
+        if isinstance(ia, ndindex.Slice) and ia.step < 0:
+            where_negative_step.append(i)
+            pos_slice = _convert_slice_with_negative_step(selection[i], x.shape[i])
+            selection[i] = pos_slice
+    where_negative_step = tuple(where_negative_step)
+    selection = tuple(selection)
 
     # Check selection is supported
-    if any(ia.step < 1 for ia in idx.args if isinstance(ia, ndindex.Slice)):
-        raise NotImplementedError(f"Slice step must be >= 1: {key}")
     if not all(
         isinstance(ia, (ndindex.Integer, ndindex.Slice, ndindex.IntegerArray))
         for ia in idx.args
@@ -140,6 +148,11 @@ def selection_function(out_key):
         if chunks != merged_chunks:
             out = merge_chunks(out, merged_chunks)
 
+    if len(where_negative_step) > 0:
+        from cubed.array_api.manipulation_functions import flip
+
+        out = flip(out, axis=where_negative_step)
+
     for axis in where_newaxis:
         from cubed.array_api.manipulation_functions import expand_dims
 
@@ -148,6 +161,16 @@ def selection_function(out_key):
     return out
 
 
+def _convert_slice_with_negative_step(key: slice, size: int) -> slice:
+    """Convert a slice with a negative step to one with a positive
+    step, which must then be followed by a flip.
+    """
+    # see https://github.com/pydata/xarray/blob/99ee8c6ca54057a9b994d7685f36236f2d5a69d9/xarray/core/indexing.py#L1056
+    start, stop, step = key.indices(size)
+    exact_stop = range(start, stop, step)[-1]
+    return slice(exact_stop, start + 1, -step)
+
+
 def _index_num_input_blocks(
     idx: ndindex.Tuple, in_chunksizes, out_chunksizes, numblocks
 ):

cubed/tests/test_array_api.py

@@ -345,6 +345,8 @@ def test_index_2d_no_op(spec, ind):
         (20, 8, slice(5, 18, 3), ((5,),)),
         # step is bigger than chunks
         (50, 5, slice(3, 50, 7), ((5, 2),)),
+        # negative step
+        (20, 4, slice(14, 3, -2), ((4, 2),)),
     ],
 )
 def test_index_1d_step(spec, shape, chunks, ind, new_chunks_expected):
@@ -364,6 +366,18 @@ def test_index_1d_step(spec, shape, chunks, ind, new_chunks_expected):
             (slice(3, 14, 2), slice(3, 14, 3)),
             ((4, 2), (3, 1),),
         ),
+        (
+            (20, 20),
+            (4, 4),
+            (slice(14, 3, -2), slice(3, 14, 3)),
+            ((4, 2), (3, 1),),
+        ),
+        (
+            (20, 20),
+            (4, 4),
+            (slice(14, 3, -2), slice(14, 3, -3)),
+            ((4, 2), (4,),),
+        ),
     ],
 )
 # fmt: on
@@ -389,12 +403,6 @@ def test_index_zero_dim(shape, chunks, ind):
     assert_array_equal(b.compute(), np.ones(shape)[ind])
 
 
-def test_index_slice_unsupported_step(spec):
-    a = xp.arange(12, chunks=(4,), spec=spec)
-    with pytest.raises(NotImplementedError):
-        a[::-1]
-
-
 @pytest.mark.parametrize("axis", [0, 1, -1, -2])
 @skip_if_cupy  # ndindex with a cupy.ndarray
 def test_take(spec, axis):