WIP: start adding top_k

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

@@ -122,14 +122,25 @@ 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"):
+    """
+    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"]')
+
+    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]
@@ -140,5 +151,5 @@ def top_k(a, k, /, axis=-1, *, mode="largest"):
 
     topk_values = np.take_along_axis(arr, topk_indices, axis=axis)
 
-    return topk_values, topk_indices
+    return Array._new(topk_values, device=a.device), Array._new(topk_indices, device=a.device)