From 5c456b26bbb3ca2089ed80b8263bfd1f0e1a642f Mon Sep 17 00:00:00 2001 From: Athan Reines Date: Wed, 1 May 2024 23:12:26 -0700 Subject: [PATCH 1/4] feat: add `count_nonzero` to specification --- .../API_specification/searching_functions.rst | 1 + .../_draft/searching_functions.py | 43 ++++++++++++++++--- 2 files changed, 37 insertions(+), 7 deletions(-) diff --git a/spec/draft/API_specification/searching_functions.rst b/spec/draft/API_specification/searching_functions.rst index c952f1aad..1a584f158 100644 --- a/spec/draft/API_specification/searching_functions.rst +++ b/spec/draft/API_specification/searching_functions.rst @@ -22,6 +22,7 @@ Objects in API argmax argmin + count_nonzero nonzero searchsorted where diff --git a/src/array_api_stubs/_draft/searching_functions.py b/src/array_api_stubs/_draft/searching_functions.py index 029459b9a..e47314e93 100644 --- a/src/array_api_stubs/_draft/searching_functions.py +++ b/src/array_api_stubs/_draft/searching_functions.py @@ -1,4 +1,4 @@ -__all__ = ["argmax", "argmin", "nonzero", "searchsorted", "where"] +__all__ = ["argmax", "argmin", "count_nonzero", "nonzero", "searchsorted", "where"] from ._types import Optional, Tuple, Literal, array @@ -54,15 +54,41 @@ def argmin(x: array, /, *, axis: Optional[int] = None, keepdims: bool = False) - """ -def nonzero(x: array, /) -> Tuple[array, ...]: +def count_nonzero( + x: array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + keepdims: bool = False, +) -> array: """ - Returns the indices of the array elements which are non-zero. + Counts the number of array elements which are non-zero. - .. note:: - If ``x`` has a complex floating-point data type, non-zero elements are those elements having at least one component (real or imaginary) which is non-zero. + Parameters + ---------- + x: array + input array. Must have a positive rank. If ``x`` is zero-dimensional, the function must raise an exception. + axis: Optional[Union[int, Tuple[int, ...]]] + axis or axes along which to count non-zero values. By default, the number of non-zero values must be computed over the entire array. If a tuple of integers, the number of non-zero values must be computed over multiple axes. Default: ``None``. + keepdims: bool + if ``True``, the reduced axes (dimensions) must be included in the result as singleton dimensions, and, accordingly, the result must be compatible with the input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes (dimensions) must not be included in the result. Default: ``False``. - .. note:: - If ``x`` has a boolean data type, non-zero elements are those elements which are equal to ``True``. + Returns + ------- + out: array + if the number of non-zeros values was computed over the entire array, a zero-dimensional array containing the total number of non-zero values; otherwise, a non-zero-dimensional array containing the counts along the specified axes. The returned array must have the default array index data type. + + Notes + ----- + + - If ``x`` has a complex floating-point data type, non-zero elements are those elements having at least one component (real or imaginary) which is non-zero. + - If ``x`` has a boolean data type, non-zero elements are those elements which are equal to ``True``. + """ + + +def nonzero(x: array, /) -> Tuple[array, ...]: + """ + Returns the indices of the array elements which are non-zero. .. admonition:: Data-dependent output shape :class: admonition important @@ -82,6 +108,9 @@ def nonzero(x: array, /) -> Tuple[array, ...]: Notes ----- + - If ``x`` has a complex floating-point data type, non-zero elements are those elements having at least one component (real or imaginary) which is non-zero. + - If ``x`` has a boolean data type, non-zero elements are those elements which are equal to ``True``. + .. versionchanged:: 2022.12 Added complex data type support. """ From 5970e6236c6fe0d0a7232585990aaa237efd6b4e Mon Sep 17 00:00:00 2001 From: Athan Reines Date: Wed, 1 May 2024 23:20:45 -0700 Subject: [PATCH 2/4] fix: address typo --- src/array_api_stubs/_draft/searching_functions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/array_api_stubs/_draft/searching_functions.py b/src/array_api_stubs/_draft/searching_functions.py index e47314e93..45bef9f89 100644 --- a/src/array_api_stubs/_draft/searching_functions.py +++ b/src/array_api_stubs/_draft/searching_functions.py @@ -102,7 +102,7 @@ def nonzero(x: array, /) -> Tuple[array, ...]: Returns ------- - out: Typle[array, ...] + out: Tuple[array, ...] a tuple of ``k`` arrays, one for each dimension of ``x`` and each of size ``n`` (where ``n`` is the total number of non-zero elements), containing the indices of the non-zero elements in that dimension. The indices must be returned in row-major, C-style order. The returned array must have the default array index data type. Notes From 991295b117793e6b9b89437e1b1c004cc4e2307d Mon Sep 17 00:00:00 2001 From: Athan Reines Date: Wed, 1 May 2024 23:26:10 -0700 Subject: [PATCH 3/4] fix: add missing import --- src/array_api_stubs/_draft/searching_functions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/array_api_stubs/_draft/searching_functions.py b/src/array_api_stubs/_draft/searching_functions.py index 45bef9f89..2ffa5f1fe 100644 --- a/src/array_api_stubs/_draft/searching_functions.py +++ b/src/array_api_stubs/_draft/searching_functions.py @@ -1,7 +1,7 @@ __all__ = ["argmax", "argmin", "count_nonzero", "nonzero", "searchsorted", "where"] -from ._types import Optional, Tuple, Literal, array +from ._types import Optional, Tuple, Literal, Union, array def argmax(x: array, /, *, axis: Optional[int] = None, keepdims: bool = False) -> array: From 14c8fbdaa2018c666bfeac64b00c13fe87e9c7e3 Mon Sep 17 00:00:00 2001 From: Athan Reines Date: Thu, 31 Oct 2024 09:58:20 -0700 Subject: [PATCH 4/4] feat!: allow zero-dimensional arrays --- src/array_api_stubs/_draft/searching_functions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/array_api_stubs/_draft/searching_functions.py b/src/array_api_stubs/_draft/searching_functions.py index 2ffa5f1fe..4eee3173b 100644 --- a/src/array_api_stubs/_draft/searching_functions.py +++ b/src/array_api_stubs/_draft/searching_functions.py @@ -67,7 +67,7 @@ def count_nonzero( Parameters ---------- x: array - input array. Must have a positive rank. If ``x`` is zero-dimensional, the function must raise an exception. + input array. axis: Optional[Union[int, Tuple[int, ...]]] axis or axes along which to count non-zero values. By default, the number of non-zero values must be computed over the entire array. If a tuple of integers, the number of non-zero values must be computed over multiple axes. Default: ``None``. keepdims: bool