diff --git a/spec/draft/API_specification/searching_functions.rst b/spec/draft/API_specification/searching_functions.rst index 1a584f158..494b72cf7 100644 --- a/spec/draft/API_specification/searching_functions.rst +++ b/spec/draft/API_specification/searching_functions.rst @@ -25,4 +25,5 @@ Objects in API count_nonzero nonzero searchsorted + top_k where diff --git a/src/array_api_stubs/_draft/searching_functions.py b/src/array_api_stubs/_draft/searching_functions.py index 291c12853..709646667 100644 --- a/src/array_api_stubs/_draft/searching_functions.py +++ b/src/array_api_stubs/_draft/searching_functions.py @@ -1,7 +1,15 @@ -__all__ = ["argmax", "argmin", "count_nonzero", "nonzero", "searchsorted", "where"] +__all__ = [ + "argmax", + "argmin", + "count_nonzero", + "nonzero", + "searchsorted", + "top_k", + "where", +] -from ._types import Optional, Tuple, Literal, Union, array +from ._types import Optional, Literal, Tuple, Union, array def argmax(x: array, /, *, axis: Optional[int] = None, keepdims: bool = False) -> array: @@ -177,6 +185,52 @@ def searchsorted( """ +def top_k( + x: array, + k: int, + /, + *, + axis: int = -1, + mode: Literal["largest", "smallest"] = "largest", +) -> Tuple[array, array]: + """ + Returns the values and indices of the ``k`` largest (or smallest) elements of an input array ``x`` along a specified dimension. + + Parameters + ---------- + x: array + input array. **Should** have a real-valued data type. + k: int + number of elements to find. **Must** be a nonnegative integer value. + axis: int + axis along which to search. A valid axis **must** be an integer on the interval ``[-N, N)``, where ``N`` is the number of axes in ``x``. If an axis is specified as a negative integer, the function **must** determine the axis along which to perform the operation by counting backward from the last axis (where ``-1`` refers to the last axis). If provided an invalid axis, the function **must** raise an exception. Default: ``-1``. + mode: Literal['largest', 'smallest'] + search mode. **Must** be one of the following modes: + + - ``'largest'``: return the ``k`` largest elements. + - ``'smallest'``: return the ``k`` smallest elements. + + Default: ``'largest'``. + + Returns + ------- + out: Tuple[array, array] + a namedtuple ``(values, indices)`` whose + + - first element **must** have the field name ``values`` and **must** be an array containing the ``k`` largest (or smallest) elements of ``x``. The array **must** have the same data type as ``x`` and **must** have the same rank (number of dimensions) and shape as ``x``, except for the axis specified by ``axis`` which **must** have size ``k``. + - second element **must** have the field name ``indices`` and **must** be an array containing indices of ``x`` that result in ``values``. The array **must** have the same shape as ``values`` and **must** have the default array index data type. + + Notes + ----- + + - If ``k`` exceeds the number of elements along the axis specified by ``axis``, behavior is left unspecified and thus implementation-dependent. Conforming implementations **may** choose, e.g., to raise an exception or return all elements. + - The order of the returned values and indices is left unspecified and thus implementation-dependent. Conforming implementations **may** return sorted or unsorted values. + - The returned indices **may** maintain the relative order of ``x`` values which compare as equal (i.e., the relative order of ``x`` values which compare as equal is implementation-dependent). + - For input arrays containing ``NaN`` values, behavior is left unspecified and thus implementation-dependent. Conforming implementations **may** choose to omit ``NaN`` values, sort ``NaN`` values to either end, or return unstable results. + - Conforming implementations **may** support complex numbers; however, inequality comparison of complex numbers is unspecified and thus implementation-dependent (see :ref:`complex-number-ordering`). + """ + + def where( condition: array, x1: Union[array, int, float, complex, bool],