Note: since argsort is part of the standard Array API, it would be possible to implement a generic yet inefficient fallback in array-api-compat while allowing to dispatch to a more efficient routine for libraries that provide it. This is what cupy.argpartition does for instance.

Thanks for the proposal @ogrisel. It's actually surprising that coverage and performance across array libraries is so spotty. I dug up the NumPy mailing list discussion, and it seemed more or less positive, just unfinished and the name to use is a nicely-sized bikeshed.

Is this function something you already have in scikit-learn internally, or are you looking for something more efficient than the argsort or similar function you use now?

In scikit-learn, for k-nearest neighbors (bruteforce exact method for medium to high dimensional space), we use a routine optimized for multicore CPUs using Cython + OpenMP for pairwise distance (similar to scipy's cdist) fused with a topk reduction implemented in templated Cython. The topk reduction itself (called "argkmin" in scikit-learn) lives here:

• https://github.com/scikit-learn/scikit-learn/blob/main/sklearn/metrics/_pairwise_distances_reduction/_argkmin.pyx.tp

This code can only be called as a reduction fused into the multithreaded pairwise distance computation kernel. It is orchestrated via:

• https://github.com/scikit-learn/scikit-learn/blob/f5ec34e0f76277ba6d0a77d3033db0af83899b64/sklearn/metrics/_pairwise_distances_reduction/_dispatcher.py#L157

For CPU, I doubt than any Array API based solution will be able to compete both on speed and memory usage.

However, we are interested in implementing Array API support for an alternative numpy code-path in order to provide GPU support, e.g. via PyTorch or CuPy. The reducer used in the numpy code-path is there:

• https://github.com/scikit-learn/scikit-learn/blob/f5ec34e0f76277ba6d0a77d3033db0af83899b64/sklearn/neighbors/_base.py#L704

It's based on numpy.argpartition followed by numpy.argsort of the top k values.

Note that to efficiently implement k-nearest neighbors in scikit-learn using the Array API, we would also need the Array API to provide scipy.spatial.distance.cdist.

I have not open an issue to discuss cdist yet. I wanted to probe the waters with topk first.

JAX also has an approximate top-k implementation specifically tuned for TPUs: https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.approx_max_k.html

I am not sure if we want to include non-exact methods in the spec. I have the feeling that there are many ways to compute such approximations and that they will require different and evolving parametrizations with different speed-accuracy trade-offs.

A PR has now been opened which proposes adding top_k and friends to the specification: #722. Please feel free to review and comment there with your concerns and feedback.