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 array indexing specification #46

Merged
merged 27 commits into from
Nov 8, 2020
Merged
Show file tree
Hide file tree
Changes from 23 commits
Commits
Show all changes
27 commits
Select commit Hold shift + click to select a range
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
158 changes: 158 additions & 0 deletions spec/API_specification/indexing.md
Original file line number Diff line number Diff line change
@@ -1 +1,159 @@
.. _indexing:

# Indexing

> Array API specification for indexing arrays.

A conforming implementation of the array API standard must adhere to the following conventions.

## Single-axis Indexing

To index a single array axis, an array must support standard Python indexing rules. Let `n` be the axis (dimension) size.

- An integer index must be an object satisfying [`operator.index`](https://www.python.org/dev/peps/pep-0357/) (e.g., `int`).

- Nonnegative indices must start at `0` (i.e., zero-based indexing).

- **Valid** nonnegative indices must reside on the half-open interval `[0, n)`.

.. note::

This specification does not explicitly require bounds checking. The behavior for out-of-bounds integer indices is left unspecified.

- Negative indices must count backward from the last array index, starting from `-1` (i.e., negative-one-based indexing, where `-1` refers to the last array index).

.. note::

A negative index `j` is equivalent to `n-j`; the former is syntactic sugar for the latter, providing a shorthand for indexing elements that would otherwise need to be specified in terms of the axis (dimension) size.

- **Valid** negative indices must reside on the closed interval `[-n, -1]`.

.. note::

This specification does not explicitly require bounds checking. The behavior for out-of-bounds integer indices is left unspecified.

- A negative index `j` is related to a zero-based nonnegative index `i` via `i = n+j`.

- Colons `:` must be used for [slices](https://docs.python.org/3/library/functions.html#slice): `start:stop:step`, where `start` is inclusive and `stop` is exclusive.

### Slice Syntax

The basic slice syntax is `i:j:k` where `i` is the starting index, `j` is the stopping index, and `k` is the step (`k != 0`). A slice may contain either one or two colons, with either an integer value or nothing on either side of each colon. The following are valid slices.

```text
A[:]
A[i:]
A[:j]
A[i:k]
A[::]
A[i::]
A[:j:]
A[::k]
A[i:j:]
A[i::k]
A[:j:k]
A[i::k]
A[i:j:k]
```

.. note::

Slice syntax can be equivalently achieved using the Python built-in [`slice()`](https://docs.python.org/3/library/functions.html#slice) API. From the perspective from `A`, the behavior of `A[i:j:k]` and `A[slice(i, j, k)]` is indistinguishable (i.e., both retrieve the same set of items from `__getitem__`).

Using a slice to index a single array axis must select `m` elements with index values

```text
i, i+k, i+2k, i+3k, ..., i+(m-1)k
```

where

```text
m = q + r
```

and `q` and `r` (`r != 0`) are the quotient and remainder obtained by dividing `j-i` by `k`

```text
j - i = qk + r
```

such that

```text
j > i + (m-1)k
```

.. note::

For `i` on the interval `[0, n)` (where `n` is the axis size), `j` on the interval `(0, n]`, `i` less than `j`, and positive step `k`, a starting index `i` is **always** included, while the stopping index `j` is **always** excluded. This preserves `x[:i]+x[i:]` always being equal to `x`.

.. note::

Using a slice to index into a single array axis should select the same elements as using a slice to index a Python list of the same size.

Slice syntax must have the following defaults. Let `n` be the axis (dimension) size.

- If `k` is not provided (e.g., `0:10`), `k` must equal `1`.
- If `k` is greater than `0` and `i` is not provided (e.g., `:10:2`), `i` must equal `0`.
- If `k` is greater than `0` and `j` is not provided (e.g., `0::2`), `j` must equal `n`.
- If `k` is less than `0` and `i` is not provided (e.g., `:10:-2`), `i` must equal `n-1`.
- If `k` is less than `0` and `j` is not provided (e.g., `0::-2`), `j` must equal `-n-1`.

Using a slice to index a single array axis must adhere to the following rules. Let `n` be the axis (dimension) size.

- If `i` equals `j`, a slice must return an empty array, whose axis (dimension) size along the indexed axis is `0`.

- If `i` and `j` resolve to starting and stopping indices, respectively, which are out-of-bounds (i.e., a starting index less than `0` and/or a stopping index greater than `n`), then the respective index is clipped to the array axis bounds. For a starting index, the bound is `0`. For a stopping index, the bound is `n`.

- Indexing a single array axis with an out-of-bounds slice (i.e., a slice which does not select any array axis elements) must return an empty array, whose axis (dimension) size along the indexed axis is `0`.

- Indexing via `:` and `::` must be equivalent and have defaults derived from the rules above. Both `:` and `::` indicate to select all elements along a single axis (dimension).

## Multi-axis Indexing

Multi-dimensional arrays must extend the concept of single-axis indexing to multiple axes by applying single-axis indexing rules along each axis (dimension) and supporting the following additional rules. Let `N` be the number of dimensions ("rank") of a multi-dimensional array `A`.

- Each axis may be independently indexed via single-axis indexing by providing a comma-separated sequence ("selection tuple") of single-axis indexing expressions (e.g., `A[:, 2:10, :, 5]`).

.. note::

In Python, `x[(exp1, exp2, ..., expN)]` is equivalent to `x[exp1, exp2, ..., expN]`; the latter is syntactic sugar for the former.

- Providing a single nonnegative integer `i` as a single-axis index must index the same elements as the slice `i:i+1`.

- Providing a single negative integer `i` as a single-axis index must index the same elements as the slice `n+i:n`, where `n` is the axis (dimension) size.

- Providing a single integer as a single-axis index must reduce the number of array dimensions by `1` (i.e., the array rank should decrease by one; if `A` has rank `2`, `rank(A)-1 == rank(A[0, :])`). In particular, a selection tuple with the `m`th element an integer (and all other entries `:`) indexes a sub-array with rank `N-1`.

- Providing a slice must retain array dimensions (i.e., the array rank must remain the same; `rank(A) == rank(A[:])`).

- For each slice which attempts to select elements along a particular axis, but whose starting index is out-of-bounds, the axis (dimension) size of the result array must be `0`.
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For devices like GPUs where you might want to allocate the memory for the result of slicing before you know the indices (as the allocation happens on the host and the indices reside on the device) this type of shape dynamism can be problematic. We should discuss what to do here in the same way as we discussed mutability.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @alextp for reviewing and raising this concern! Would definitely be good to discuss. cc @rgommers

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would this also be an issue for slices bounds that are beyond the array shape (e.g., a[0:1000] where a has shape (100,))?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think slice clipping could be omitted from the spec. In other words, the spec only specifies what happens for slices where the start or stop are in [-size, size] (where the ends of that interval may or may not be included depending on the different cases of start/stop and negative/positive step).

Slices on Python lists implement "clipping" behavior to the size, which NumPy matches, but this could be something that isn't specified in the spec, and libraries could do something else. It is also possible to "manually" implement clipping in user code, in much a similar way that you can "manually" implement bounds checking (clipping happens first when a slice is computed, see the a[-100::2] example here).

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both slices and array shapes are usually kept on the host, so I think you should be able to figure out the necessary buffer sizes without any synchronization. Or am I missing something?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@apaszke reading the first comment from @alextp, I think he meant something like

ix_start = x.sum()  # computed index depends on values in GPU array `x`
z = y[ix_start : ix_start + 2]  # z will have size 2, unless ix_start is >= y.shape[0] - 1

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've removed any bounds checking requirements from the proposal, under the rationale that how an array implementation handles out-of-bounds indices (including within slices) is best left to the implementation.

Based on your example @rgommers and @alextp's comment, the gist of the problem is that I may want to allocate memory for an array having axis size 2, regardless of whether indices are valid (in-bounds) or not. And if we required that, in this case, out-of-bounds indices unconditionally must result in an array whose axis size is 0, then I won't be able to allocate memory without knowing the actual index values, thus triggering device syncs and triggering a perf cliff.


- If the number of provided single-axis indexing expressions is less than `N`, then `:` must be assumed for the remaining dimensions (e.g., if `A` has rank `2`, `A[2:10] == A[2:10, :]`).

- An `IndexError` exception must be raised if the number of provided single-axis indexing expressions is greater than `N`.

- Providing [ellipsis](https://docs.python.org/3/library/constants.html#Ellipsis) must apply `:` to each dimension necessary to index all dimensions (e.g., if `A` has rank `4`, `A[1:, ..., 2:5] == A[1:, :, :, 2:5]`). Only a single ellipsis must be allowed. An `IndexError` exception must be raised if more than one ellipsis is provided.

- The result of multi-axis indexing must be an array of the same data type as the indexed array.

## Boolean Array Indexing

An array must support indexing via a **single** `M`-dimensional boolean array `B` with shape `S1 = (s1, ..., sM)` according to the following rules. Let `A` be an `N`-dimensional array with shape `S2 = (s1, ..., sM, ..., sN)`.

- If `N >= M`, then `A[B]` must replace the first `M` dimensions of `A` with a single dimension having a size equal to the number of `True` elements in `B`.

.. note::

For example, if `N == M == 2`, indexing `A` via a boolean array `B` will return a one-dimensional array whose size is equal to the number of `True` elements in `B`.

- If `N < M`, then an `IndexError` exception must be raised.

- The size of each dimension in `B` must equal the size of the corresponding dimension in `A` or be `0`, beginning with the first dimension in `A`. If a dimension size does not equal the size of the corresponding dimension in `A` and is not `0`, then an `IndexError` exception must be raised.

- The elements of a boolean index array must be iterated in row-major, C-style order, with the exception of zero-dimensional boolean arrays.

- A zero-dimensional boolean index array (equivalent to `True` or `False`) must follow the same axis replacement rules stated above. Namely, a zero-dimensional boolean index array removes zero dimensions and adds a single dimension of length `1` if the index array's value is `True` and of length `0` if the index array's value is `False`. Accordingly, for a zero-dimensional boolean index array `B`, the result of `A[B]` has shape `S = (1, s1, ..., sN)` if the index array's value is `True` and has shape `S = (0, s1, ..., sN)` if the index array's value is `False`.

- The result of indexing into an array via a boolean index array must be an array of the same data type as the indexed array.
2 changes: 2 additions & 0 deletions spec/API_specification/searching_functions.md
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
.. _searching-functions:

# Searching Functions

> Array API specification for functions for searching arrays.
Expand Down