ENH: add top_k

[ev-br]

Add basic top_k wrappers, per data-apis/array-api#722
-tests tracker: data-apis/array-api-tests#438

• torch
• numpy : works locally with ENH: Add top_k based on partitioning numpy/numpy#31659
• cupy
• dask

Supersedes and closes gh-158

To test locally: with this PR and data-apis/array-api-tests#438, use

$ ARRAY_API_TESTS_SKIP_DTYPES=uint16,uint32,uint64 ARRAY_API_TESTS_VERSION=draft ARRAY_API_TESTS_MODULE=array_api_compat.torch pytest array_api_tests/test_searching_functions.py::test_top_k -s --max-examples=5000

[Copilot]

Pull request overview

Adds an Array API-compatible top_k wrapper to the array_api_compat.torch namespace, mapping the standardized signature to torch.topk.

Changes:

• Introduces top_k(a, k, *, axis, mode, **kwargs) wrapper in the Torch backend.
• Validates mode and forwards arguments to torch.topk.
• Exports top_k via the Torch backend __all__.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[ev-br]

So torch.topk needs only a small tweak to make it spec-compliant, which is exactly the remit of array-api-compat.
numpy is getting a native implementation in 2.6 (?).

Next is what to do about older numpies and cupy, which do not have native implementations.
We technically could reuse the full argsort-based implementation, cf data-apis/array-api-strict#215, but it has a wrong big-O complexity, and I'd be loathe sneak it in.
So unless there are very strong opinions to the contrary, I'd leave them out of array-api-compat. Thoughts?

[rgommers]

Older numpy/cupy: I can see the case either way. Pinging @betatim for scikit-learn needs, and @MaanasArora in case he has good ideas after working on the numpy implementation.

[betatim]

What are the options?

My two thoughts/inputs:

• for scikit-learn we say that our array API support requires you to use the latest versions of all the dependencies. Older versions are "untested, might work or not". (I need to think about how this will work if we enable it the support by default)
• having an implementation for every library in array-api-compat is what I'd expect. Even if that implementation is slower/less optimal than it could be. array-api-compat is a "functionality provider of last resort" - ideally top_k is in the actual library but if not it is better to have it here than not at all. The alternative is that we write our own implementation in scikit-learn which would suffer from the same problems as the one here but is more work for scikit-learn and benefits less people

[MaanasArora]

We technically could reuse the full argsort-based implementation

As far as I understand, for NumPy at least (and the link you referenced) we are using argpartition instead of argsorts. This is the same operation that will be used in the native implementation for NumPy 2.6. The missing feature in numpy <= 2.5 is descending argpartitions, which means that the largest=True case requires special Python-level handling to flip the kth and ensure nans are handled correctly (which was present in the older NumPy implementation that data-apis/array-api-strict#215 seems to use).

Given that the most expensive underlying operation is the same, I wouldn't expect an outsized performance difference for a compatibility layer. Whether the complexity of manual nan and kth handling is worth supporting older NumPy is a different question. Given the performance should be in a similar order, I'd lean towards supporting if it is very useful for client libraries.