From 31ca6cb43cc597b4205dacce2f400fb18d50343c Mon Sep 17 00:00:00 2001 From: Athan Reines Date: Sun, 5 Nov 2023 19:14:38 -0800 Subject: [PATCH 01/10] feat: add specification for `searchsorted` --- .../API_specification/searching_functions.rst | 1 + .../_draft/searching_functions.py | 29 +++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/spec/draft/API_specification/searching_functions.rst b/spec/draft/API_specification/searching_functions.rst index 01ab4e82a..c952f1aad 100644 --- a/spec/draft/API_specification/searching_functions.rst +++ b/spec/draft/API_specification/searching_functions.rst @@ -23,4 +23,5 @@ Objects in API argmax argmin nonzero + searchsorted where diff --git a/src/array_api_stubs/_draft/searching_functions.py b/src/array_api_stubs/_draft/searching_functions.py index e586a7656..89c5b74ae 100644 --- a/src/array_api_stubs/_draft/searching_functions.py +++ b/src/array_api_stubs/_draft/searching_functions.py @@ -87,6 +87,35 @@ def nonzero(x: array, /) -> Tuple[array, ...]: """ +def searchsorted( + x1: array, + x2: array, + /, + *, + side: Literal["left", "right"] = "left", + sorter: Optional[array] = None, +) -> array: + """ + Finds the indices into ``x1`` such that, if the corresponding elements in ``x2`` were inserted before the indices, the order of ``x1``, when sorted in ascending order, would be preserved. + + Parameters + ---------- + x1: array + input array. Must be a one-dimensional array. If ``sorter`` is ``None``, must be sorted in ascending order; otherwise, ``sorter`` must be an array of indices that sort ``x1`` in ascending order. + x2: array + array containing search values. + side: Literal['left', 'right'] + if ``'left'``, then each returned index ``i`` must satisfy ``x1[i-1] < v[j] <= x1[i]``; otherwise, if ``'right'``, the each returned index ``i`` must satisfy ``x1[i-1] <= v[j] < x1[i]``. Default: ``'left'``. + sorter: Optional[array] + array of indices that sort ``x1`` in ascending order. The array must be one-dimensional and have an integer data type. Default: ``None``. + + Returns + ------- + out: array + an array of indices with the same shape as ``x2``. The returned array must have the default array index data type. + """ + + def where(condition: array, x1: array, x2: array, /) -> array: """ Returns elements chosen from ``x1`` or ``x2`` depending on ``condition``. From 7a217c81485078346b037414cc356c0a896f74e4 Mon Sep 17 00:00:00 2001 From: Athan Reines Date: Sun, 5 Nov 2023 19:33:38 -0800 Subject: [PATCH 02/10] fix: correct the required `sorter` shape and clarify edge cases --- src/array_api_stubs/_draft/searching_functions.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/array_api_stubs/_draft/searching_functions.py b/src/array_api_stubs/_draft/searching_functions.py index 89c5b74ae..8e23e768a 100644 --- a/src/array_api_stubs/_draft/searching_functions.py +++ b/src/array_api_stubs/_draft/searching_functions.py @@ -105,9 +105,9 @@ def searchsorted( x2: array array containing search values. side: Literal['left', 'right'] - if ``'left'``, then each returned index ``i`` must satisfy ``x1[i-1] < v[j] <= x1[i]``; otherwise, if ``'right'``, the each returned index ``i`` must satisfy ``x1[i-1] <= v[j] < x1[i]``. Default: ``'left'``. + if ``'left'``, then each returned index ``i`` must satisfy ``x1[i-1] < x2[j] <= x1[i]``, and, if, for an element ``x2[j]``, no index satisfies the index condition, then the returned index for that element must be ``0``. Otherwise, if ``'right'``, then each returned index ``i`` must satisfy ``x1[i-1] <= x2[j] < x1[i]``, and, if, for an element ``x2[j]``, no index satisfies the index condition, then the returned index for that element must be ``N``, where ``N`` is the number of elements in ``x1``. Default: ``'left'``. sorter: Optional[array] - array of indices that sort ``x1`` in ascending order. The array must be one-dimensional and have an integer data type. Default: ``None``. + array of indices that sort ``x1`` in ascending order. The array must have the same shape as ``x1`` and have an integer data type. Default: ``None``. Returns ------- From 1ac2ea71ab00e6275e770ebdaff0df7303050ea2 Mon Sep 17 00:00:00 2001 From: Athan Reines Date: Sun, 5 Nov 2023 19:35:58 -0800 Subject: [PATCH 03/10] fix: use multi-dimensional index syntax --- 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 8e23e768a..b8fe36dcc 100644 --- a/src/array_api_stubs/_draft/searching_functions.py +++ b/src/array_api_stubs/_draft/searching_functions.py @@ -105,7 +105,7 @@ def searchsorted( x2: array array containing search values. side: Literal['left', 'right'] - if ``'left'``, then each returned index ``i`` must satisfy ``x1[i-1] < x2[j] <= x1[i]``, and, if, for an element ``x2[j]``, no index satisfies the index condition, then the returned index for that element must be ``0``. Otherwise, if ``'right'``, then each returned index ``i`` must satisfy ``x1[i-1] <= x2[j] < x1[i]``, and, if, for an element ``x2[j]``, no index satisfies the index condition, then the returned index for that element must be ``N``, where ``N`` is the number of elements in ``x1``. Default: ``'left'``. + if ``'left'``, then each returned index ``i`` must satisfy ``x1[i-1] < x2[m][n]...[j] <= x1[i]``, and, if, for an element ``x2[m][n]...[j]``, no index satisfies the index condition, then the returned index for that element must be ``0``. Otherwise, if ``'right'``, then each returned index ``i`` must satisfy ``x1[i-1] <= x2[m][n]...[j] < x1[i]``, and, if, for an element ``x2[m][n]...[j]``, no index satisfies the index condition, then the returned index for that element must be ``N``, where ``N`` is the number of elements in ``x1``. Default: ``'left'``. sorter: Optional[array] array of indices that sort ``x1`` in ascending order. The array must have the same shape as ``x1`` and have an integer data type. Default: ``None``. From b89ae7de04825701e23eb66ad25bd86b2081a2a6 Mon Sep 17 00:00:00 2001 From: Athan Reines Date: Sun, 5 Nov 2023 19:37:22 -0800 Subject: [PATCH 04/10] fix: import `Literal` type --- 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 b8fe36dcc..5ddac9c98 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", "nonzero", "where"] -from ._types import Optional, Tuple, array +from ._types import Optional, Tuple, Literal, array def argmax(x: array, /, *, axis: Optional[int] = None, keepdims: bool = False) -> array: From c0594bf03fac4fa60776dd363d34d711b7376c60 Mon Sep 17 00:00:00 2001 From: Athan Reines Date: Sun, 5 Nov 2023 19:46:06 -0800 Subject: [PATCH 05/10] fix: add `searchsorted` to list of exported functions --- 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 5ddac9c98..d04711823 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", "where"] +__all__ = ["argmax", "argmin", "nonzero", "searchsort", "where"] from ._types import Optional, Tuple, Literal, array From 0a47a677f866692829e7b352d214dff03cc30b16 Mon Sep 17 00:00:00 2001 From: Athan Reines Date: Sun, 5 Nov 2023 19:49:54 -0800 Subject: [PATCH 06/10] fix: use correct API name --- 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 d04711823..1ea186284 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", "searchsort", "where"] +__all__ = ["argmax", "argmin", "nonzero", "searchsorted", "where"] from ._types import Optional, Tuple, Literal, array From 13165dd4df37d0b0177a9a0204849ff09d4b59ae Mon Sep 17 00:00:00 2001 From: Athan Reines Date: Wed, 29 Nov 2023 23:17:58 -0800 Subject: [PATCH 07/10] Update copy for increased clarity --- .../_draft/searching_functions.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/src/array_api_stubs/_draft/searching_functions.py b/src/array_api_stubs/_draft/searching_functions.py index 1ea186284..20df824f9 100644 --- a/src/array_api_stubs/_draft/searching_functions.py +++ b/src/array_api_stubs/_draft/searching_functions.py @@ -105,7 +105,21 @@ def searchsorted( x2: array array containing search values. side: Literal['left', 'right'] - if ``'left'``, then each returned index ``i`` must satisfy ``x1[i-1] < x2[m][n]...[j] <= x1[i]``, and, if, for an element ``x2[m][n]...[j]``, no index satisfies the index condition, then the returned index for that element must be ``0``. Otherwise, if ``'right'``, then each returned index ``i`` must satisfy ``x1[i-1] <= x2[m][n]...[j] < x1[i]``, and, if, for an element ``x2[m][n]...[j]``, no index satisfies the index condition, then the returned index for that element must be ``N``, where ``N`` is the number of elements in ``x1``. Default: ``'left'``. + argument controlling which index is returned if a value lands exactly on an edge. + + Let ``x`` be an array of rank ``N`` where ``v`` is an individual element given by ``v = x2[n,m,...,j]``. + + If ``side == 'left'``, then + + - each returned index ``i`` must satisfy the index condition ``x1[i-1] < v <= x1[i]``. + - if no index satisfies the index condition, then the returned index for that element must be ``0``. + + Otherwise, if ``side == 'right'``, then + + - each returned index ``i`` must satisfy the index condition ``x1[i-1] <= v < x1[i]``. + - if no index satisfies the index condition, then the returned index for that element must be ``N``, where ``N`` is the number of elements in ``x1``. + + Default: ``'left'``. sorter: Optional[array] array of indices that sort ``x1`` in ascending order. The array must have the same shape as ``x1`` and have an integer data type. Default: ``None``. From 79da2ea7ed5e1cbce711e75c2e4f81bc25343762 Mon Sep 17 00:00:00 2001 From: Athan Reines Date: Wed, 10 Jan 2024 23:53:04 -0800 Subject: [PATCH 08/10] Add note concerning sort order and behavior when an array contains NaNs --- src/array_api_stubs/_draft/searching_functions.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/array_api_stubs/_draft/searching_functions.py b/src/array_api_stubs/_draft/searching_functions.py index 20df824f9..2e2462005 100644 --- a/src/array_api_stubs/_draft/searching_functions.py +++ b/src/array_api_stubs/_draft/searching_functions.py @@ -101,9 +101,9 @@ def searchsorted( Parameters ---------- x1: array - input array. Must be a one-dimensional array. If ``sorter`` is ``None``, must be sorted in ascending order; otherwise, ``sorter`` must be an array of indices that sort ``x1`` in ascending order. + input array. Must be a one-dimensional array. Should have a real-valued data type. If ``sorter`` is ``None``, must be sorted in ascending order; otherwise, ``sorter`` must be an array of indices that sort ``x1`` in ascending order. x2: array - array containing search values. + array containing search values. Should have a real-valued data type. side: Literal['left', 'right'] argument controlling which index is returned if a value lands exactly on an edge. @@ -127,6 +127,13 @@ def searchsorted( ------- out: array an array of indices with the same shape as ``x2``. The returned array must have the default array index data type. + + Notes + ----- + + For real-valued floating-point arrays, the sort order of NaNs and signed zeros is unspecified and thus implementation-dependent. Accordingly, when a real-valued floating-point array contains NaNs and signed zeros, what constitutes ascending order may vary among specification-conforming array libraries. + + Specification-conforming libraries should, however, ensure consistency with ``sort`` and ``argsort`` (i.e., if a value in ``x2`` is inserted into ``x1`` according to the corresponding index in the output array and ``sort`` is invoked on the resultant array, the sorted result should be an array in the same order). """ From 73878c897ded3b82d99a1a483d2f0ed28f0ab57e Mon Sep 17 00:00:00 2001 From: Athan Reines Date: Wed, 10 Jan 2024 23:56:55 -0800 Subject: [PATCH 09/10] Update note --- 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 2e2462005..d84debdda 100644 --- a/src/array_api_stubs/_draft/searching_functions.py +++ b/src/array_api_stubs/_draft/searching_functions.py @@ -131,7 +131,7 @@ def searchsorted( Notes ----- - For real-valued floating-point arrays, the sort order of NaNs and signed zeros is unspecified and thus implementation-dependent. Accordingly, when a real-valued floating-point array contains NaNs and signed zeros, what constitutes ascending order may vary among specification-conforming array libraries. + For real-valued floating-point arrays, the sort order of NaNs and signed zeros is unspecified and thus implementation-dependent. Accordingly, when a real-valued floating-point array contains NaNs and signed zeros, what constitutes ascending order may vary among specification-conforming array libraries and behavior is implementation-dependent. Specification-conforming libraries should, however, ensure consistency with ``sort`` and ``argsort`` (i.e., if a value in ``x2`` is inserted into ``x1`` according to the corresponding index in the output array and ``sort`` is invoked on the resultant array, the sorted result should be an array in the same order). """ From b89a6cdf7c346325c9dff6ff604b287679cbc72b Mon Sep 17 00:00:00 2001 From: Athan Reines Date: Thu, 11 Jan 2024 00:00:13 -0800 Subject: [PATCH 10/10] Update note --- src/array_api_stubs/_draft/searching_functions.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/array_api_stubs/_draft/searching_functions.py b/src/array_api_stubs/_draft/searching_functions.py index d84debdda..dda000e74 100644 --- a/src/array_api_stubs/_draft/searching_functions.py +++ b/src/array_api_stubs/_draft/searching_functions.py @@ -131,9 +131,9 @@ def searchsorted( Notes ----- - For real-valued floating-point arrays, the sort order of NaNs and signed zeros is unspecified and thus implementation-dependent. Accordingly, when a real-valued floating-point array contains NaNs and signed zeros, what constitutes ascending order may vary among specification-conforming array libraries and behavior is implementation-dependent. + For real-valued floating-point arrays, the sort order of NaNs and signed zeros is unspecified and thus implementation-dependent. Accordingly, when a real-valued floating-point array contains NaNs and signed zeros, what constitutes ascending order may vary among specification-conforming array libraries. - Specification-conforming libraries should, however, ensure consistency with ``sort`` and ``argsort`` (i.e., if a value in ``x2`` is inserted into ``x1`` according to the corresponding index in the output array and ``sort`` is invoked on the resultant array, the sorted result should be an array in the same order). + While behavior for arrays containing NaNs and signed zeros is implementation-dependent, specification-conforming libraries should, however, ensure consistency with ``sort`` and ``argsort`` (i.e., if a value in ``x2`` is inserted into ``x1`` according to the corresponding index in the output array and ``sort`` is invoked on the resultant array, the sorted result should be an array in the same order). """