RFC: assignment via integer array indexing

[crusaderky]

[EDIT] this issue is about the lack of xp.put or equivalent __setitem_ semantics for integar array indices. Read comments below.

[original post]
The current situation when an index is an array of ints or bools is very messy:

• numpy blindly passes every unexpected object it finds to np.asarray, which makes it accept np.ndarray, lists, tuples, but also memoryviews and anything with an __array__ interface
• Sparse accepts lists, tuples, or numpy arrays; other sparse arrays don't work
• PyTorch won't accept numpy arrays of unsigned integers
• PyTorch won't accept PyTorch arrays of integers with dtype other than the native int
• JAX won't accept lists or tuples
• dask won't accept tuples

I think the array API standard should define a reasonable common surface.

xref data-apis/array-api-compat#205 (comment)

NOTE: the much more complicated case of multiple fancy indices, e.g. a[[0, 1], [1, 2]] is out of scope for this issue.

[rgommers]

Just a quick note that take is the standardized/uniform way of doing indexing with an integer array (xref gh-416 and the issue that PR links to).

You're completely right that it's very messy across libraries - it's just very hard, since we can't mandate that libraries stop doing something and break backwards compatible (or technically we can, but it'll probably just be ignored rather than implemented).

[crusaderky]

Just a quick note that take is the standardized/uniform way of doing indexing with an integer array (xref gh-416 and the issue that PR links to).

I couldn't find put or anything similar though? How do you update along a fancy index?

[ev-br]

cross-ref #900 which aims to actually standardize it.

[ev-br]

Looks like gh-900 covered this request, and we can close it as completed, I'd think.

[kgryte]

@ev-br #900 doesn't cover setitem semantics, however.

[ev-br]

Would be great to update the issue title/test then, if we repurpose it for __setitem__ specifically.

[amacati]

Reviving this issue as we are currently facing the limitation in scipy/scipy#23425. Is there any movement on this? The case for __setitem__ semantics with integer arrays seems both sensible and widely implemented. When running the current implementation in scipy/scipy#23425, the only framework that fails to set items is array-api-strict.

[rgommers]

Is there any movement on this?

No, no change, sorry.

the only framework that fails to set items is array-api-strict.

I don't think that's possible, JAX doesn't implement __setitem__. It's either not tested there or the SciPy implementation does an out of place update I suspect.

[amacati]

Ah right, then maybe this is a topic for array-api-extra instead. Because this pattern should be common in most frameworks:

import array_api_strict as xp
import jax.numpy as jp

x = jp.array([[1, 2, 3], [4, 5, 6]])
idx = jp.array([0])
# This could be emulated in array-api-extra with xpx.at(x)[idx, ...] = xp.asarray([[7, 8, 9]])
# It is already emulating this case for regular indexing
new_x = x.at[idx, ...].set(jp.array([[7, 8, 9]]))
print(new_x)  # [[7 8 9], [4 5 6]]

x = xp.asarray([[1, 2, 3], [4, 5, 6]])
idx = xp.asarray([0])
new_x = x[idx, ...] = xp.asarray([[7, 8, 9]])  # IndexError: Fancy indexing __setitem__ is not supported.

[rgommers]

Ah right, then maybe this is a topic for array-api-extra instead.

Indeed: https://data-apis.org/array-api-extra/generated/array_api_extra.at.html

[amacati]

Using the current spec, is there an array API-compatible way to implement this feature though? Because swapping the last line with

import array_api_extra as xpx
...
new_x = xpx.at(x)[idx, ...].set(xp.asarray([[7, 8, 9]]))  # IndexError: ...

raises the same error, since array-api-extra tries to translate this into x[idx] = y for array-api-strict.

[crusaderky]

Using the current spec, is there an array API-compatible way to implement this feature though? Because swapping the last line with

import array_api_extra as xpx
...
new_x = xpx.at(x)[idx, ...].set(xp.asarray([[7, 8, 9]])) # IndexError: ...
raises the same error, since array-api-extra tries to translate this into x[idx] = y for array-api-strict.

It is possible to update xpx.at to detect an integer array index and force an out-of-place operation in that case.