WIP: update top_k to match the spec proposal

Author: ev-br

array_api_strict/__init__.py

@@ -297,9 +297,9 @@
 
 __all__ += ["concat", "expand_dims", "flip", "moveaxis", "permute_dims", "repeat", "reshape", "roll", "squeeze", "stack", "tile", "unstack"]
 
-from ._searching_functions import argmax, argmin, nonzero, count_nonzero, searchsorted, where
+from ._searching_functions import argmax, argmin, nonzero, count_nonzero, searchsorted, where, top_k
 
-__all__ += ["argmax", "argmin", "nonzero", "count_nonzero", "searchsorted", "where"]
+__all__ += ["argmax", "argmin", "nonzero", "count_nonzero", "searchsorted", "where", "top_k"]
 
 from ._set_functions import unique_all, unique_counts, unique_inverse, unique_values, isin
 

array_api_strict/_searching_functions.py

@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import Literal, NamedTuple
 
 import numpy as np
 
@@ -9,6 +9,11 @@
 from ._helpers import _maybe_normalize_py_scalars
 
 
+class TopKResult(NamedTuple):
+    values: Array
+    indices: Array
+
+
 def argmax(x: Array, /, *, axis: int | None = None, keepdims: bool = False) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.argmax <numpy.argmax>`.
@@ -122,23 +127,46 @@ def where(condition: Array, x1: Array | complex, x2: Array | complex, /) -> Arra
     return Array._new(np.where(condition._array, x1._array, x2._array), device=x1.device)
 
 
-def top_k(a, k, /, axis=-1, *, mode="largest"):
+
+def top_k(
+    a: Array,
+    k: int,
+    /,
+    *,
+    axis: int =-1,
+    mode: Literal["largest", "smallest"] = "largest"
+) -> TopKResult:
+    """
+    Array API compatible wrapper for :py:func:`np.top_k <numpy.top_k>`.
+
+    See its docstring for more information.
+    """
     if k <= 0:
         raise ValueError(f'k(={k}) provided must be positive.')
 
-    positive_axis = axis if axis > 0 else axis % arr.ndim
+    if mode not in ["largest", "smallest"]:
+        raise ValueError(f'{mode = } not in ["largest", "smallest"]')
+
+    if k > a.shape[axis]:
+        raise ValueError(f"{k = } exceeds {a.shape[axis] = }") 
+
+    positive_axis = axis if axis > 0 else axis % a.ndim
+
+    arr = a._array
 
     slice_start = (np.s_[:],) * positive_axis
-    if largest:
+    if mode == "largest":
         indices_array = np.argpartition(arr, -k, axis=axis)
-        slice = slice_start + (np.s_[-k:],)
-        topk_indices = indices_array[slice]
+        slice_ = slice_start + (np.s_[-k:],)
+        topk_indices = indices_array[slice_]
     else:
         indices_array = np.argpartition(arr, k-1, axis=axis)
-        slice = slice_start + (np.s_[:k],)
-        topk_indices = indices_array[slice]
+        slice_ = slice_start + (np.s_[:k],)
+        topk_indices = indices_array[slice_]
 
     topk_values = np.take_along_axis(arr, topk_indices, axis=axis)
 
-    return topk_values, topk_indices
-
+    return TopKResult(
+        Array._new(topk_values, device=a.device),
+        Array._new(topk_indices, device=a.device)
+    )