Skip to content

Commit df8081c

Browse files
authored
docs: clarify support for negative indices in take and take_along_axis
PR-URL: #894 Reviewed-by: Matt Haberland
1 parent 071780c commit df8081c

File tree

1 file changed

+11
-15
lines changed

1 file changed

+11
-15
lines changed

src/array_api_stubs/_draft/indexing_functions.py

+11-15
Original file line numberDiff line numberDiff line change
@@ -7,21 +7,14 @@ def take(x: array, indices: array, /, *, axis: Optional[int] = None) -> array:
77
"""
88
Returns elements of an array along an axis.
99
10-
.. note::
11-
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.
12-
1310
Parameters
1411
----------
1512
x: array
1613
input array. Should have one or more dimensions (axes).
1714
indices: array
18-
array indices. The array must be one-dimensional and have an integer data type.
19-
20-
.. note::
21-
This specification does not require bounds checking. The behavior for out-of-bounds indices is left unspecified.
22-
15+
array indices. The array must be one-dimensional and have an integer data type. If an index is negative, the function must determine the element to select along a specified axis (dimension) by counting from the last element (where ``-1`` refers to the last element).
2316
axis: Optional[int]
24-
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.
17+
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 (where ``-1`` refers to the last dimension).
2518
2619
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.
2720
@@ -33,6 +26,8 @@ def take(x: array, indices: array, /, *, axis: Optional[int] = None) -> array:
3326
Notes
3427
-----
3528
29+
- 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.
30+
- This specification does not require bounds checking. The behavior for out-of-bounds indices is left unspecified.
3631
- When ``x`` is a zero-dimensional array, behavior is unspecified and thus implementation-defined.
3732
3833
.. versionadded:: 2022.12
@@ -51,16 +46,17 @@ def take_along_axis(x: array, indices: array, /, *, axis: int = -1) -> array:
5146
x: array
5247
input array. Must be compatible with ``indices``, except for the axis (dimension) specified by ``axis`` (see :ref:`broadcasting`).
5348
indices: array
54-
array indices. Must have the same rank (i.e., number of dimensions) as ``x``.
55-
56-
.. note::
57-
This specification does not require bounds checking. The behavior for out-of-bounds indices is left unspecified.
58-
49+
array indices. Must have the same rank (i.e., number of dimensions) as ``x``. If an index is negative, the function must determine the element to select along a specified axis (dimension) by counting from the last element (where ``-1`` refers to the last element).
5950
axis: int
60-
axis along 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. Default: ``-1``.
51+
axis along 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 (where ``-1`` refers to the last dimension). Default: ``-1``.
6152
6253
Returns
6354
-------
6455
out: array
6556
an array having the same data type as ``x``. Must have the same rank (i.e., number of dimensions) as ``x`` and must have a shape determined according to :ref:`broadcasting`, except for the axis (dimension) specified by ``axis`` whose size must equal the size of the corresponding axis (dimension) in ``indices``.
57+
58+
Notes
59+
-----
60+
61+
- This specification does not require bounds checking. The behavior for out-of-bounds indices is left unspecified.
6662
"""

0 commit comments

Comments
 (0)