feat: add API specification for returning the k largest elements

spec/draft/API_specification/searching_functions.rst

@@ -25,4 +25,5 @@ Objects in API
    count_nonzero
    nonzero
    searchsorted
+   top_k
    where

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],