RFC: add APIs for setting elements via an array of indices (i.e., put, put_along_axis, etc)

Open
#979 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
25/100
Issue type
Feature
Clarity
Needs clarification
Activity status
Stale
Tech stack
numpy, python, pytorch, tensorflow
Domain
data

Research direction

Start with the array API indexing discussion in gh-177 and the linked array-api issue comment, then compare the listed NumPy, CuPy, JAX, PyTorch, TensorFlow, Dask, and Ndonnx APIs. Done means the proposal defines a settled, interoperable API and semantics for setting elements by indices, including behavior for immutable arrays.

Written by the indexing model from the issue text.

Description

Copied and adapted from @kgryte's proposal at gh-177 per https://github.com/data-apis/array-api/issues/177#issuecomment-2883011323

Proposal

Add APIs for setting elements via an array of indices.

Motivation

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

The main argument is that Indexing does not currently support providing a array of indices to index into an array. The principal reason for not supporting fancy indexing is . However, use of fancy indexing is relatively common in NumPy and similar libraries where dynamically setting rows/cols/values is possible and can be readily implemented. Another reason for not support fancy indexing mutation is that some libraries feature immutable arrays. array_api_extra.at.set demonstrates that it is often sufficient to mutate where possible and create a copy with the specified updates otherwise.

Background

The following table summarizes library implementations of such APIs:

op NumPy CuPy JAX.numpy Torch Tensorflow Dask.array Ndonnx
setting elements along axis put put put scatter_? scatter_nd? ?
setting elements over matching 1d slices put_along_axis put_along_axis put_along_axis scatter_? See pytorch/pytorch#120209. scatter_nd? ?

Discussion about these function in gh-177 concluded with https://github.com/data-apis/array-api/issues/177#issuecomment-1514155595, especially:

The JAX issue is most difficult to resolve (can be done, but a lot of work still to deal with read-only views or similar), but the lack of API uniformity makes this a hard sell in general.

This seems to be resolved. jax.numpy.put and jax.numpy.put_along_axis are implemented; they just return modifyied copies rather than mutating the array in place.

Dominant language
Python
Stars
281
Forks
52
No merged PRs in 30d

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from data-apis/array-api

All issues in data-apis/array-api

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.