forked from data-apis/array-api
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathset_functions.py
More file actions
175 lines (118 loc) · 11.6 KB
/
set_functions.py
File metadata and controls
175 lines (118 loc) · 11.6 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
__all__ = ["unique_all", "unique_counts", "unique_inverse", "unique_values"]
from ._types import Tuple, array
def unique_all(x: array, /) -> Tuple[array, array, array, array]:
"""
Returns the unique elements of an input array ``x``, the first occurring indices for each unique element in ``x``, the indices from the set of unique elements that reconstruct ``x``, and the corresponding counts for each unique element in ``x``.
.. admonition:: Data-dependent output shape
:class: important
The shapes of two of the 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, et cetera) can 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.
Parameters
----------
x: array
input array. If ``x`` has more than one dimension, the function **must** flatten ``x`` and return the unique elements of the flattened array.
Returns
-------
out: Tuple[array, array, array, array]
a namedtuple ``(values, indices, inverse_indices, counts)`` whose
- first element **must** have the field name ``values`` and **must** be a one-dimensional array containing the unique elements of ``x``. The array **must** have the same data type as ``x``.
- second element **must** have the field name ``indices`` and **must** be an array containing the indices (first occurrences) of a flattened ``x`` that result in ``values``. The array **must** have the same shape as ``values`` and **must** have the default array index data type.
- third element **must** have the field name ``inverse_indices`` and **must** be an array containing the indices of ``values`` that reconstruct ``x``. The array **must** have the same shape as ``x`` and **must** have the default array index data type.
- fourth element **must** have the field name ``counts`` and **must** be an array containing the number of times each unique element occurs in ``x``. The order of the returned counts **must** match the order of ``values``, such that a specific element in ``counts`` corresponds to the respective unique element in ``values``. The returned array **must** have same shape as ``values`` and **must** have the default array index data type.
Notes
-----
- The order of unique elements returned by this function is unspecified and thus implementation-defined. As a consequence, element order **may** vary between implementations.
- Uniqueness **should** be determined based on value equality (see :func:`~array_api.equal`). For input arrays having floating-point data types, value-based equality implies the following behavior.
- As ``nan`` values compare as ``False``, ``nan`` values **should** be considered distinct.
- As complex floating-point values having at least one ``nan`` component compare as ``False``, complex floating-point values having ``nan`` components **should** be considered distinct.
- As ``-0`` and ``+0`` compare as ``True``, signed zeros **should not** be considered distinct, and the corresponding unique element **may** be implementation-defined (e.g., an implementation **may** choose to return ``-0`` if ``-0`` occurs before ``+0``).
As signed zeros are not distinct, using ``inverse_indices`` to reconstruct the input array is not guaranteed to return an array having the exact same values.
Each ``nan`` value and each complex floating-point value having a ``nan`` component **should** have a count of one, while the counts for signed zeros **should** be aggregated as a single count.
.. versionchanged:: 2022.12
Added complex data type support.
.. versionchanged:: 2023.12
Clarified flattening behavior and required the order of ``counts`` match the order of ``values``.
"""
def unique_counts(x: array, /) -> Tuple[array, array]:
"""
Returns the unique elements of an input array ``x`` and the corresponding counts for each unique element in ``x``.
.. admonition:: Data-dependent output shape
:class: important
The shapes of two of the 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.) can 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.
Parameters
----------
x: array
input array. If ``x`` has more than one dimension, the function **must** flatten ``x`` and return the unique elements of the flattened array.
Returns
-------
out: Tuple[array, array]
a namedtuple `(values, counts)` whose
- first element **must** have the field name ``values`` and **must** be a one-dimensional array containing the unique elements of ``x``. The array **must** have the same data type as ``x``.
- second element **must** have the field name `counts` and **must** be an array containing the number of times each unique element occurs in ``x``. The order of the returned counts **must** match the order of ``values``, such that a specific element in ``counts`` corresponds to the respective unique element in ``values``. The returned array **must** have same shape as ``values`` and **must** have the default array index data type.
Notes
-----
- The order of unique elements returned by this function is unspecified and thus implementation-defined. As a consequence, element order **may** vary between implementations.
- Uniqueness **should** be determined based on value equality (see :func:`~array_api.equal`). For input arrays having floating-point data types, value-based equality implies the following behavior.
- As ``nan`` values compare as ``False``, ``nan`` values **should** be considered distinct.
- As complex floating-point values having at least one ``nan`` component compare as ``False``, complex floating-point values having ``nan`` components **should** be considered distinct.
- As ``-0`` and ``+0`` compare as ``True``, signed zeros **should not** be considered distinct, and the corresponding unique element **may** be implementation-defined (e.g., an implementation **may** choose to return ``-0`` if ``-0`` occurs before ``+0``).
Each ``nan`` value and each complex floating-point value having a ``nan`` component **should** have a count of one, while the counts for signed zeros **should** be aggregated as a single count.
.. versionchanged:: 2022.12
Added complex data type support.
.. versionchanged:: 2023.12
Clarified flattening behavior and required the order of ``counts`` match the order of ``values``.
"""
def unique_inverse(x: array, /) -> Tuple[array, array]:
"""
Returns the unique elements of an input array ``x`` and the indices from the set of unique elements that reconstruct ``x``.
.. admonition:: Data-dependent output shape
:class: important
The shapes of two of the 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.) can 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.
Parameters
----------
x: array
input array. If ``x`` has more than one dimension, the function **must** flatten ``x`` and return the unique elements of the flattened array.
Returns
-------
out: Tuple[array, array]
a namedtuple ``(values, inverse_indices)`` whose
- first element **must** have the field name ``values`` and **must** be a one-dimensional array containing the unique elements of ``x``. The array **must** have the same data type as ``x``.
- second element **must** have the field name ``inverse_indices`` and **must** be an array containing the indices of ``values`` that reconstruct ``x``. The array **must** have the same shape as ``x`` and have the default array index data type.
Notes
-----
- The order of unique elements returned by this function is unspecified and thus implementation-defined. As a consequence, element order **may** vary between implementations.
- Uniqueness **should** be determined based on value equality (see :func:`~array_api.equal`). For input arrays having floating-point data types, value-based equality implies the following behavior.
- As ``nan`` values compare as ``False``, ``nan`` values **should** be considered distinct.
- As complex floating-point values having at least one ``nan`` component compare as ``False``, complex floating-point values having ``nan`` components **should** be considered distinct.
- As ``-0`` and ``+0`` compare as ``True``, signed zeros **should not** be considered distinct, and the corresponding unique element **may** be implementation-defined (e.g., an implementation **may** choose to return ``-0`` if ``-0`` occurs before ``+0``).
As signed zeros are not distinct, using ``inverse_indices`` to reconstruct the input array is not guaranteed to return an array having the exact same values.
.. versionchanged:: 2022.12
Added complex data type support.
.. versionchanged:: 2023.12
Clarified flattening behavior.
"""
def unique_values(x: array, /) -> array:
"""
Returns the unique elements of an input array ``x``.
.. admonition:: Data-dependent output shape
:class: important
The shapes of two of the 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.) can 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.
Parameters
----------
x: array
input array. If ``x`` has more than one dimension, the function **must** flatten ``x`` and return the unique elements of the flattened array.
Returns
-------
out: array
a one-dimensional array containing the set of unique elements in ``x``. The returned array **must** have the same data type as ``x``.
Notes
-----
- The order of unique elements returned by this function is unspecified and thus implementation-defined. As a consequence, element order **may** vary between implementations.
- Uniqueness **should** be determined based on value equality (see :func:`~array_api.equal`). For input arrays having floating-point data types, value-based equality implies the following behavior.
- As ``nan`` values compare as ``False``, ``nan`` values **should** be considered distinct.
- As complex floating-point values having at least one ``nan`` component compare as ``False``, complex floating-point values having ``nan`` components **should** be considered distinct.
- As ``-0`` and ``+0`` compare as ``True``, signed zeros **should not** be considered distinct, and the corresponding unique element **may** be implementation-defined (e.g., an implementation **may** choose to return ``-0`` if ``-0`` occurs before ``+0``).
.. versionchanged:: 2022.12
Added complex data type support.
.. versionchanged:: 2023.12
Required that the output array must be one-dimensional.
"""