Repository navigation
ENH: Adding partition and argpartition?聽#448
Description
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! 馃檹
Activity
cakedev0 commentedon Sep 30, 2025
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 commentedon Oct 1, 2025
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 commentedon Oct 1, 2025
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 commentedon Oct 1, 2025
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