Skip to content

Commit b7b0dcd

Browse files
kgrytesteff456rgommersshoyerleofang
authored
Clarify guidance on operations having data-dependent output shapes (#193)
* Add note for functions that doesn't support boolean array indexing * clarify the notes * remove note from where function * Create a unique admonition for data-dependent shapes * Rewrite note with the given suggestions * Update notes * Add document stub * Add admonition for boolean array indexing and update reference * Add explainer on data-dependent output shapes * Fix typo * Update copy Co-authored-by: Stephan Hoyer <shoyer@google.com> * Update copy Co-authored-by: Leo Fang <leofang@bnl.gov> * Update copy Co-authored-by: Stephannie Jimenez <steff456@hotmail.com> Co-authored-by: Ralf Gommers <ralf.gommers@gmail.com> Co-authored-by: Stephan Hoyer <shoyer@google.com> Co-authored-by: Leo Fang <leofang@bnl.gov>
1 parent 8dade11 commit b7b0dcd

File tree

5 files changed

+34
-0
lines changed

5 files changed

+34
-0
lines changed

spec/API_specification/indexing.md

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,12 @@ _Rationale: this is consistent with bounds-checking for single-axis indexing. An
160160

161161
## Boolean Array Indexing
162162

163+
:::{admonition} Data-dependent output shape
164+
:class: important
165+
166+
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.
167+
:::
168+
163169
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)`.
164170

165171
- 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)]`.

spec/API_specification/searching_functions.md

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,12 @@ Returns the indices of the minimum values along a specified axis. When the minim
6969
(function-nonzero)=
7070
### nonzero(x, /)
7171

72+
:::{admonition} Data-dependent output shape
73+
:class: important
74+
75+
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.
76+
:::
77+
7278
Returns the indices of the array elements which are non-zero.
7379

7480
#### Parameters

spec/API_specification/set_functions.md

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,12 @@ A conforming implementation of the array API standard must provide and support t
1515
(function-unique)=
1616
### unique(x, /, *, return_counts=False, return_index=False, return_inverse=False)
1717

18+
:::{admonition} Data-dependent output shape
19+
:class: important
20+
21+
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.
22+
:::
23+
1824
Returns the unique elements of an input array `x`.
1925

2026
#### Parameters
Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
(data-dependent-output-shapes)=
2+
3+
# Data-dependent output shapes
4+
5+
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.
6+
7+
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.
8+
9+
Value-dependent operations are demarcated in this specification using an admonition similar to the following:
10+
11+
:::{admonition} Data-dependent output shape
12+
:class: important
13+
14+
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.
15+
:::

spec/design_topics/index.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ Design topics & constraints
66
:maxdepth: 1
77

88
copies_views_and_mutation
9+
data_dependent_output_shapes
910
data_interchange
1011
device_support
1112
static_typing

0 commit comments

Comments
 (0)