ENH: add partition and argpartition functions

[cakedev0]

Closes#448

Supports nd-arrays.

The crux of this PR is the torch part: transforming kthvalue outputs to a partition/argpartition. The rest is mostly wrappers and checks.

[lucascolley]

I don't know how to handle the sparse backend (is_pydata_sparse_namespace): it has no argsort...

see

array-api-extra/tests/test_funcs.py

Line 1183 in ca20f03

@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no argsort")

[lucascolley]

thanks @cakedev0 !

[lucascolley]

thanks @cakedev0 !

Looks like there is also a merge conflict now.

[lucascolley]

hmmm, I would rather not rely on undocumented behaviour. Is there an alternative?

[cakedev0]

Fair ^^

Three options:

• add an assert a_left.max() == a_left[k]
• We can just re-run the same logic with kth=1 and largest=True. Impact on perfs is probably 10 to 100% slower depending on the input. But it doens't add a lot of logic
• We can do a if a_left.max() != a_left[k]: swap_max_with_last_element(a_left, axis=-1) => requires to implement swap_max_with_last_element (and the equivalent for argsort).

I vote for 1 because I'm lazy but I like perf :p

[cakedev0]

Edit: wait I need to rethink something about numpy.partition specs...

[cakedev0]

So! I rewrote entirely this section, it now relies on torch.kthvalue and is very aligned with numpy's behavior.

On a side note: the description of the behavior of the partition function in numpy is fairly blurry when the k-th element has duplicates... In practice, numpy does a tree-way partitioning: <, == and >. I reproduced this behavior in my new torch implementation, but jax doesn't (I tried to test the tree-way partitioning and jax fails it...).

I will maybe open an issue on numpy to ask for some clarification.

[lucascolley]

Good idea!

[ogrisel]

On a side note: the description of the behavior of the partition function in numpy is fairly blurry when the k-th element has duplicates... In practice, numpy does a tree-way partitioning: <, == and >. I reproduced this behavior in my new torch implementation, but jax doesn't (I tried to test the tree-way partitioning and jax fails it...).

It might be worth contributing this consideration to the array API spec discussion:

• RFC issue: RFC: add topk and / or argpartition array-api#629
• SPEC updated PR: feat: add API specification for returning the k largest elements array-api#722

[lucascolley]

thanks @cakedev0, looks close!

[cakedev0]

Thanks for the reactive and helpful reviews @lucascolley

Sorry for the numpy docs style details, I'll make sure to read the doc about this carefully before opening another PR 😉

[lucascolley]

thanks @cakedev0, let's merge it! And thanks for taking a look too @ogrisel .

Would be great to follow-up on #449 (comment).

Are you interested in taking over gh-341 next?

[ogrisel]

@lucascolley I see that this as been milestoned for 0.9.1, but those are new functions. Wouln'd it make sense to make them part of a new 0.10.0 release instead?

[lucascolley]

@lucascolley I see that this as been milestoned for 0.9.1, but those are new functions. Wouln'd it make sense to make them part of a new 0.10.0 release instead?

we use https://jacobtomlinson.dev/effver/, so no. Unless you see any problems with that!

[ogrisel]

Ok, as you wish.