Skip to content

Add take specification for returning elements of an array along a specified axis #416

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 10 commits into from
Nov 16, 2022
1 change: 1 addition & 0 deletions spec/API_specification/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ API specification
elementwise_functions
function_and_method_signatures
indexing
indexing_functions
linear_algebra_functions
manipulation_functions
searching_functions
Expand Down
28 changes: 28 additions & 0 deletions spec/API_specification/indexing_functions.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
.. _indexing-functions:

Indexing Functions
===================

Array API specification for functions for indexing arrays.

A conforming implementation of the array API standard must provide and support the following functions adhering to the following conventions.

- Positional parameters must be `positional-only <https://www.python.org/dev/peps/pep-0570/>`_ parameters. Positional-only parameters have no externally-usable name. When a function accepting positional-only parameters is called, positional arguments are mapped to these parameters based solely on their order.
- Optional parameters must be `keyword-only <https://www.python.org/dev/peps/pep-3102/>`_ arguments.
- Broadcasting semantics must follow the semantics defined in :ref:`broadcasting`.
- Unless stated otherwise, functions must support the data types defined in :ref:`data-types`.
- Unless stated otherwise, functions must adhere to the type promotion rules defined in :ref:`type-promotion`.

Objects in API
--------------

.. currentmodule:: signatures.indexing_functions

..
NOTE: please keep the functions in alphabetical order

.. autosummary::
:toctree: generated
:template: method.rst

take
28 changes: 28 additions & 0 deletions spec/API_specification/signatures/indexing_functions.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
from ._types import Union, array

def take(x: array, indices: Union[int, array], /, *, axis: int) -> array:
"""
Returns elements of an array along an axis.

.. note::
Conceptually, the call ``take(x, indices, axis=3)`` is equivalent to ``x[:,:,:,indices,...]``; however, explicit indexing via arrays of indices is not currently supported in this specification due to concerns regarding ``__setitem__`` and array mutation semantics.

Parameters
----------
x: array
input array.
indices: Union[int, array]
array indices. If ``indices`` is an integer (or a zero-dimensional array satisfying ``operator.index``), the function must follow multi-axis indexing semantics when providing an integer for a single-axis index (see :ref:`indexing`). If ``indices`` is a non-zero-dimensional array, the array must be one-dimensional and have an integer data type.
axis: int
axis over which to select values. If ``axis`` is negative, the function must determine the axis along which to select values by counting from the last dimension.

Returns
-------
out: array
an array having the same data type as ``x``. If ``indices`` is an array, the output array must have the same rank (i.e., number of dimensions) as ``x`` and must have the same shape as ``x``, except for the axis specified by ``axis`` whose size must equal the number of elements in ``indices``. If ``indices`` is an integer, the output array must have one less dimension than ``x`` (i.e., the array rank should decrease by one). In particular, if ``x`` has rank ``N``, when ``indices`` is an integer, this function must be equivalent to a selection tuple (e.g., ``x[:,:,3,:]``, ``x[0,:]``, et cetera) with the ``m``\th element an integer (and all other entries ``:``), thus indexing a sub-array with rank ``N-1``.

.. note::
When ``indices`` is an integer and ``x`` is an array of rank ``1``, the returned array must be an array of rank ``0``, not a scalar.
"""

__all__ = ['take']