Skip to content
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

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

Merged
merged 10 commits into from
Nov 16, 2022
1 change: 1 addition & 0 deletions spec/API_specification/array_api/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from .data_type_functions import *
import array_api.data_types as dtype
from .elementwise_functions import *
from .indexing_functions import *
from .linear_algebra_functions import *
from .manipulation_functions import *
from .searching_functions import *
Expand Down
27 changes: 27 additions & 0 deletions spec/API_specification/array_api/indexing_functions.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
from ._types import Union, array

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

.. note::
Conceptually, ``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: array
array indices. 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.

If ``x`` is a one-dimensional array, providing an ``axis`` is optional; however, if ``x`` has more than one dimension, providing an ``axis`` is required.

Returns
-------
out: array
an array having the same data type as ``x``. 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``.
"""

__all__ = ['take']
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:: array_api

..
NOTE: please keep the functions in alphabetical order

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

take