Skip to content

Clarify guidance on operations having data-dependent output shapes #193

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 14 commits into from
Jun 28, 2021
Merged
6 changes: 6 additions & 0 deletions spec/API_specification/indexing.md
Original file line number Diff line number Diff line change
Expand Up @@ -160,6 +160,12 @@ _Rationale: this is consistent with bounds-checking for single-axis indexing. An

## Boolean Array Indexing

:::{admonition} Data-dependent output shape
:class: important

For common boolean array use cases (e.g., using a dynamically-sized boolean array mask to filter the values of another array), the shape of the output array is data-dependent; hence, array libraries which build computation graphs (e.g., JAX, Dask, etc.) may find boolean array indexing difficult to implement. Accordingly, such libraries may choose to omit boolean array indexing. See {ref}`data-dependent-output-shapes` section for more details.
:::

An array must support indexing where the **sole index** is an `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`. The values in the resulting array must be in row-major (C-style order); this is equivalent to `A[nonzero(B)]`.
Expand Down
6 changes: 6 additions & 0 deletions spec/API_specification/searching_functions.md
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,12 @@ Returns the indices of the minimum values along a specified axis. When the minim
(function-nonzero)=
### nonzero(x, /)

:::{admonition} Data-dependent output shape
:class: important

The shape of the output array for this function depends on the data values in the input array; hence, array libraries which build computation graphs (e.g., JAX, Dask, etc.) may find this function difficult to implement without knowing array values. Accordingly, such libraries may choose to omit this function. See {ref}`data-dependent-output-shapes` section for more details.
:::

Returns the indices of the array elements which are non-zero.

#### Parameters
Expand Down
6 changes: 6 additions & 0 deletions spec/API_specification/set_functions.md
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,12 @@ A conforming implementation of the array API standard must provide and support t
(function-unique)=
### unique(x, /, *, return_counts=False, return_index=False, return_inverse=False)

:::{admonition} Data-dependent output shape
:class: important

The shapes of one or more of output arrays for this function depend on the data values in the input array; hence, array libraries which build computation graphs (e.g., JAX, Dask, etc.) may find this function difficult to implement without knowing array values. Accordingly, such libraries may choose to omit this function. See {ref}`data-dependent-output-shapes` section for more details.
:::

Returns the unique elements of an input array `x`.

#### Parameters
Expand Down
15 changes: 15 additions & 0 deletions spec/design_topics/data_dependent_output_shapes.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
(data-dependent-output-shapes)=

# Data-dependent output shapes

Array libraries which build computation graphs commonly employ static analysis that relies upon known shapes. For example, JAX requires known array sizes when compiling code, in order to perform static memory allocation. Functions and operations which are value-dependent present difficulties for such libraries, as array sizes cannot be inferred ahead of time without also knowing the contents of the respective arrays.

While value-dependent functions and operations are not impossible to implement for array libraries which build computation graphs, this specification does not want to impose an undue burden on such libraries and permits omission of value-dependent operations. All other array libraries are expected, however, to implement the value-dependent operations included in this specification in order to be array specification compliant.

Value-dependent operations are demarcated in this specification using an admonition similar to the following:

:::{admonition} Data-dependent output shape
:class: important

The shape of the output array for this function/operation depends on the data values in the input array; hence, array libraries which build computation graphs (e.g., JAX, Dask, etc.) may find this function/operation difficult to implement without knowing array values. Accordingly, such libraries may choose to omit this function. See {ref}`data-dependent-output-shapes` section for more details.
:::
1 change: 1 addition & 0 deletions spec/design_topics/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ Design topics & constraints
:maxdepth: 1

copies_views_and_mutation
data_dependent_output_shapes
data_interchange
device_support
static_typing
Expand Down