Add array indexing specification

[kgryte]

This PR

• specifies array indexing semantics.

Notes

• bounds checking: JAX/Numba) doesn't seem to perform bounds checks for perf reasons (see jax.numpy array indexing has different out-of-bounds behavior to numpy jax-ml/jax#278). Out-of-bounds index behavior left unspecified.

• Left out np.newaxis in favor of adding expand_dims to specification.

• Apart from boolean array indexing, NumPy's "advanced indexing" is not included and is not universally among array libraries (e.g., dask and TensorFlow do not support all of NumPy's indexing semantics).

[alextp]

For devices like GPUs where you might want to allocate the memory for the result of slicing before you know the indices (as the allocation happens on the host and the indices reside on the device) this type of shape dynamism can be problematic. We should discuss what to do here in the same way as we discussed mutability.

[kgryte]

Thanks @alextp for reviewing and raising this concern! Would definitely be good to discuss. cc @rgommers

[asmeurer]

Would this also be an issue for slices bounds that are beyond the array shape (e.g., a[0:1000] where a has shape (100,))?

[asmeurer]

I think slice clipping could be omitted from the spec. In other words, the spec only specifies what happens for slices where the start or stop are in [-size, size] (where the ends of that interval may or may not be included depending on the different cases of start/stop and negative/positive step).

Slices on Python lists implement "clipping" behavior to the size, which NumPy matches, but this could be something that isn't specified in the spec, and libraries could do something else. It is also possible to "manually" implement clipping in user code, in much a similar way that you can "manually" implement bounds checking (clipping happens first when a slice is computed, see the a[-100::2] example here).

[apaszke]

Both slices and array shapes are usually kept on the host, so I think you should be able to figure out the necessary buffer sizes without any synchronization. Or am I missing something?

[rgommers]

@apaszke reading the first comment from @alextp, I think he meant something like

ix_start = x.sum()  # computed index depends on values in GPU array `x`
z = y[ix_start : ix_start + 2]  # z will have size 2, unless ix_start is >= y.shape[0] - 1

[kgryte]

I've removed any bounds checking requirements from the proposal, under the rationale that how an array implementation handles out-of-bounds indices (including within slices) is best left to the implementation.

Based on your example @rgommers and @alextp's comment, the gist of the problem is that I may want to allocate memory for an array having axis size 2, regardless of whether indices are valid (in-bounds) or not. And if we required that, in this case, out-of-bounds indices unconditionally must result in an array whose axis size is 0, then I won't be able to allocate memory without knowing the actual index values, thus triggering device syncs and triggering a perf cliff.

[kgryte]

@rgommers This should be ready for review given removal of bounds checking requirements.

[rgommers]

Okay, let's get this in. The rationale for lack of specifying clipping looks fine, there may be follow-up discussion - we can expand that paragraph if needed.

[rgommers]

Thanks for the input everyone!