ENH: Adding partition and argpartition?

[cakedev0]

I'd like to add partition and argpartition in this lib, that would help me in one of my PR in scikit-learn: scikit-learn/scikit-learn#32288 😄

I'd love to do it myself, but I need some guidance. Here is what I want to do:

• jax/numpy/cupy implement those functions, so just call them.
• in torch you have torch.topk that can be used to implement those functions in O(n) (instead of O(n log n) if you rely on the sort). Same in TensorFlow, you have tf.math.top_k. In Dask too: dask.array.topk
• the default case would be to use sort/argsort

My main question is: how should I implement the control-flow "if torch: do this; if dask: do that; [...]"?

Thank you! 🙏

[cakedev0]

I looked around and I think adding those function would be really worth-it. It would allow a much more efficient implementation of the quantile function in SciPy and here: #341

Indeed, sorting is not needed to compute quantiles, only partitioning, and partitioning is O(n) 😄

[lucascolley]

Thanks @cakedev0 !

The array-agnostic implementation would live in https://github.com/data-apis/array-api-extra/blob/main/src/array_api_extra/_lib/_funcs.py. The delegation to existing libraries would live in https://github.com/data-apis/array-api-extra/blob/main/src/array_api_extra/_delegation.py. See also our contributing docs: https://data-apis.org/array-api-extra/contributing.html.

Would you like to make a PR?

[cakedev0]

Ah I hadn't see _delegation.py, now that I see how the code looks like in this file, it's pretty clear what I need to do. I'd like to make the PR yes!

[cakedev0]

Ok, thanks again for the pointers @lucascolley, writing the PR went rather smoothly 👌 Can you take a quick look at the PR to check it goes in the right direction? #449

And also I'd like to have your opinion about what to do with the sparse backend (see the PR desc).

Then, I'll add nd-support and polish it for the real review :)

Edit: ah you're too reactive, you already reviewed my PR lol