test indexing with arrays

[ev-br]

cross-ref data-apis/array-api#669 : test indexing arrays objects with a mix of integers and arrays.

• I'm not sure how to properly assert correctness here. I currently just check that results agree with numpy, which is not right. Maybe there's something in the ndindex library that would help. Could you weight in @asmeurer?
• Do we need to test __setitem__ in addition to getitem?

As a data point, the current (hacky) test seems to pass locally with numpy (unsurprisingly), pytorch and jax.numpy.
EDIT: also passes locally with array-api-strict, patched via https://github.com/data-apis/array-api-strict/compare/main...ev-br:array_indexing?expand=1 (the patch certainly allows way too much, is for demo purposes only).

A CLI incantation to run tests locally

$ ARRAY_API_TESTS_MODULE=torch pytest array_api_tests/test_array_object.py::test_getitem_arrays_and_ints -vs --max-examples=2000 --pdb

cc @kgryte for viz

[asmeurer]

I think the main way to test correctness is to assert each element corresponds to indexing the individual element of the index array. This would make the test unvectorized, though. For a vectorized test the only thing I can think of is to flatten things and test that it matches take.

As far as mixed indices, ndindex can help by canonicalizing the index (you want the expand() method).

[ev-br]

Run the (current, hacky) test on dask: it does not seem to allow array indexing.

$ ARRAY_API_TESTS_MODULE=array_api_compat.dask.array pytest array_api_tests/test_array_object.py::test_getitem_arrays_and_ints -vs --max-examples=2000 --pdb
...

??? x.shape = (1,) key = [dask.array<reshape, shape=(1, 1), dtype=int32, chunksize=(1, 1), chunktype=numpy.ndarray>] -- [(1, 1)]
FAILED
>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> traceback >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>

array_api_tests/test_array_object.py:286: in test_getitem_arrays_and_ints
    out = x[key]
../../miniforge3/envs/array-api-tests/lib/python3.11/site-packages/dask/array/core.py:2006: in __getitem__
    self, index2 = slice_with_int_dask_array(self, index2)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

x = dask.array<array, shape=(1,), dtype=int32, chunksize=(1,), chunktype=numpy.ndarray>
index = (dask.array<reshape, shape=(1, 1), dtype=int32, chunksize=(1, 1), chunktype=numpy.ndarray>,)

...

            if isinstance(idx, Array) and idx.dtype.kind in "iu":
                if idx.ndim == 0:
                    idx = idx[np.newaxis]
                    x = slice_with_int_dask_array_on_axis(x, idx, out_axis)
                    x = x[tuple(0 if i == out_axis else slice(None) for i in range(x.ndim))]
                    dropped_axis_cnt += 1
                elif idx.ndim == 1:
                    x = slice_with_int_dask_array_on_axis(x, idx, out_axis)
                    out_index.append(slice(None))
                else:
>                   raise NotImplementedError(
                        "Slicing with dask.array of ints only permitted when "
                        "the indexer has zero or one dimensions"
                    )
E                   NotImplementedError: Slicing with dask.array of ints only permitted when the indexer has zero or one dimensions
E                   Falsifying example: test_getitem_arrays_and_ints(
E                       shape=(1,),
E                       data=data(...),
E                   )
E                   Draw 1 (obj): [0]
E                   Draw 2: True
E                   Draw 3: ((1, 1),)
E                   Draw 4: dask.array<reshape, shape=(1, 1), dtype=int32, chunksize=(1, 1), chunktype=numpy.ndarray>
E                   Explanation:
E                       These lines were always and only run by failing examples:
E                           /home/br/miniforge3/envs/array-api-tests/lib/python3.11/site-packages/dask/array/reshape.py:344
E                           /home/br/miniforge3/envs/array-api-tests/lib/python3.11/site-packages/dask/array/reshape.py:346
E                           /home/br/miniforge3/envs/array-api-tests/lib/python3.11/site-packages/dask/array/reshape.py:348
E                           /home/br/miniforge3/envs/array-api-tests/lib/python3.11/site-packages/dask/array/reshape.py:349
E                           /home/br/miniforge3/envs/array-api-tests/lib/python3.11/site-packages/dask/array/reshape.py:350
E                           (and 4 more with settings.verbosity >= verbose)

../../miniforge3/envs/array-api-tests/lib/python3.11/site-packages/dask/array/slicing.py:960: NotImplementedError

[ev-br]

I think the main way to test correctness is to assert each element corresponds to indexing the individual element of the index array. This would make the test unvectorized, though. For a vectorized test the only thing I can think of is to flatten things and test that it matches take.

Thanks @asmeurer !
I actually wonder about a small generalization of what's currently in this PR: convert to numpy instead of an making it unvectorized. Will need a small utility to paper over downstream libs, but that's simple.
This is clearly a hack, even if a pragmatic one. I'm sure you have considered something like this previously @asmeurer ?

[ev-br]

cross-ref #343 for a cleaner but slightly more restricted test (which only generates 1D indexing arrays)

[ev-br]

closing in favor of #343