Skip to content

ENH: Adding partition and argpartition?聽#448

Closed
Feature
#449
Closed
ENH: Adding partition and argpartition?#448
Feature
#449
Labels
enhancementNew feature or requestnew function

Description

@cakedev0
cakedev0
opened on Sep 30, 2025

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

  1. cakedev0 commented on Sep 30, 2025

    @cakedev0
    cakedev0
    on Sep 30, 2025
    ContributorAuthor

    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) 馃槃

  2. lucascolley commented on Oct 1, 2025

    @lucascolley
    lucascolley
    on Oct 1, 2025
    Member

    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?

  3. cakedev0 commented on Oct 1, 2025

    @cakedev0
    cakedev0
    on Oct 1, 2025
    ContributorAuthor

    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!

  4. lucascolley
    added
    enhancementNew feature or request
    new function
    on Oct 1, 2025
  5. lucascolley
    added the
    Feature
    issue type on Oct 1, 2025

  6. cakedev0 commented on Oct 1, 2025

    @cakedev0
    cakedev0
    on Oct 1, 2025
    ContributorAuthor

    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

  7. lucascolley
    closed this as completedin #449on Oct 3, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or requestnew function

    Type

    Feature

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions