From c822d5892059d2f1dc52247b9d15631b532eb1c0 Mon Sep 17 00:00:00 2001 From: Evgeni Burovski Date: Wed, 1 Jul 2026 12:05:33 +0200 Subject: [PATCH 1/2] ENH: add a basic top_k implementation Taken from https://github.com/data-apis/array-api-compat/pull/158/ Co-authored-by: Jules --- array_api_strict/_searching_functions.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/array_api_strict/_searching_functions.py b/array_api_strict/_searching_functions.py index 6b8659b..d1449da 100644 --- a/array_api_strict/_searching_functions.py +++ b/array_api_strict/_searching_functions.py @@ -123,3 +123,25 @@ def where(condition: Array, x1: Array | complex, x2: Array | complex, /) -> Arra x1, x2 = Array._normalize_two_args(x1, x2) return Array._new(np.where(condition._array, x1._array, x2._array), device=x1.device) + + +def top_k(a, k, /, axis=-1, *, mode="largest"): + if k <= 0: + raise ValueError(f'k(={k}) provided must be positive.') + + positive_axis = axis if axis > 0 else axis % arr.ndim + + slice_start = (np.s_[:],) * positive_axis + if largest: + indices_array = np.argpartition(arr, -k, axis=axis) + 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] + + topk_values = np.take_along_axis(arr, topk_indices, axis=axis) + + return topk_values, topk_indices + From 6925ca96e05e3d117ea6f99979f90330e2887f13 Mon Sep 17 00:00:00 2001 From: Evgeni Burovski Date: Wed, 1 Jul 2026 12:46:25 +0200 Subject: [PATCH 2/2] WIP: update top_k to match the spec proposal --- array_api_strict/__init__.py | 3 +- array_api_strict/_searching_functions.py | 48 +++++++++++++++++++----- 2 files changed, 40 insertions(+), 11 deletions(-) diff --git a/array_api_strict/__init__.py b/array_api_strict/__init__.py index d1763a4..e5ee294 100644 --- a/array_api_strict/__init__.py +++ b/array_api_strict/__init__.py @@ -303,10 +303,11 @@ count_nonzero, nonzero, searchsorted, + top_k, where, ) -__all__ += ["argmax", "argmin", "count_nonzero", "nonzero", "searchsorted", "where"] +__all__ += ["argmax", "argmin", "count_nonzero", "nonzero", "searchsorted", "top_k", "where"] from ._set_functions import ( isin, diff --git a/array_api_strict/_searching_functions.py b/array_api_strict/_searching_functions.py index d1449da..5796d7e 100644 --- a/array_api_strict/_searching_functions.py +++ b/array_api_strict/_searching_functions.py @@ -1,4 +1,4 @@ -from typing import Literal +from typing import Literal, NamedTuple import numpy as np @@ -12,6 +12,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 `. @@ -125,23 +130,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 `. + + 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) + )