Proposal: add APIs for getting and setting elements via a list of indices (i.e., `take`, `put`, etc)

Proposal

Add APIs for getting and setting elements via a list of indices.

Motivation

Currently, the array API specification does not provide a direct means of extracting and setting a list of elements along an axis. Such operations are relatively common in NumPy usage either via “fancy indexing” or via explicit take and put APIs.

Two main arguments come to mind for supporting at least basic take and put APIs:

1. Indexing does not currently support providing a list of indices to index into an array. The principal reason for not supporting fancy indexing stems from dynamic shapes and compatibility with accelerator libraries. However, use of fancy indexing is relatively common in NumPy and similar libraries where dynamically extracting rows/cols/values is possible and can be readily implemented.

2. Currently, the output of a subset of APIs currently included in the standard cannot be readily consumed without manual workarounds if a specification-conforming library implemented only the APIs in the standard. For example,

   • argsort returns an array of indices. In NumPy, the output of this function can be consumed by put_along_axis and take_along_axis.
   • unique can return an array of indices if return_index is True.

Background

The following table summarizes library implementations of such APIs:

op	NumPy	CuPy	Dask	MXNet	Torch	TensorFlow
extracting elements along axis	take	take	take	take	take/gather	gather/numpy.take
setting elements along axis	put	put	–	–	scatter	scatter_nd/tensor_scatter_nd_update
extracting elements over matching 1d slices	take_along_axis	take_along_axis	–	–	–	gather_nd/numpy.take_alongaxis
setting elements over matching 1d slices	put_along_axis	–	–	–	–	–

While most libraries implement some form of take, fewer implement other complementary APIs.

[rgommers]

Support for take has been merged, see gh-416.

[rgommers]

Thanks for the ping on this issue @arogozhnikov - and nice to see the experimental work on indexing in einops. I’d like to see gh-416 finished and merged in the coming days to indeed add take support with 1-D indices.