RFC: add APIs for getting elements via a list of indices (i.e., `take`, `take_along_axis`, etc)

Open
#177 38 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
Domain
api, data

Research direction

Start with the proposal and the linked sorting_functions.md and set_functions.md specification sections, then compare the referenced NumPy, CuPy, Dask, MXNet, Torch, and TensorFlow APIs. The work is done when the project agrees on the scope and semantics of the take/put family and the relevant array API specification is updated.

Written by the indexing model from the issue text.

Description

API extension topic: Dynamic Shapes

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.

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.