Add numpy argpartition function support by rragundez · Pull Request #8732 · numba/numba

rragundez · 2023-01-28T06:58:40Z

Add numpy argpartition function support.

I used some of the implementation of partition. Something to note it that I modified the already existing _partition_factory to add an optional argument if the user would need np.argpartition instead of np.partition, as you will see in the code it doesn't look the most pretty, but it was either that or write a _argpartition_factory function which has exactly the same logic it just swaps index elements at the same time it swaps the array with the values being sorted.

Let me know if you prefer to leave it as is or go the duplication route. I think this is the best compromise to not repeat code and have a single source of truth of partition logic.

Solves #2445
#4074 would need to be updated.

cheers

kc611

Hi @rragundez , Thank you for the PR 😄 . Your contributions to Numba are highly appreciated.

I gave the PR an initial review and it mostly LGTM!

A few points to note:

The logic does seem a bit off within the _partition_factory/_partition, but as you mentioned and from a maintainer's point of view, it's less code to manage and reuses existing logic so it does makes sense to keep it this way.
The np.argpartition, based on the existing np.partition logic does seem to deviate from NumPy behaviour, they are implemented using different algorithms (see #3320 (comment) for original np.partition implementation discussion) and whilst this might seem a bit implicit, since the ordering of the elements in the two partitions is supposed to be undefined, but it might be a good idea to mention the behaviour in the docs that the outputs may deviate from NumPy.

kc611

A small suggestion regarding default function arguments.

On this note is there any particular reason you went with the default argument for idx being np.arange(0) ?

rragundez · 2023-01-30T19:35:24Z

@kc611 Thanks for your review and feedback.
I also do not like it at all, but since the function is passed to register_jitable it needs to make sense of the types of the variables and if left with default as None it cannot make sense of the None type and the array of type int64 which I or idx can become. So I decided to use the np.arange(0) as default value which is an empty array of type int64.
Does that make sense? Let me know what you think.

rragundez · 2023-02-06T18:36:08Z

@kc611 Hi, did you have time to check my remark? perhaps I am wrong but that behaviour I describe is what I noticed when developing the PR.

kc611 · 2023-02-07T03:59:54Z

Hi @rragundez, apologies for the delay. Following a small internal discussion, it'd be better for us to go with None as default argument. The reason being; having mutable default arguments is generally not recommended for Python functions due to the kind of unexpected behaviours it leads to, down the line.

Additionally, the type inference within Numba is expected to figure out if the argument type is None and prune the branches accordingly during compile time.

rragundez · 2023-02-07T12:44:54Z

@kc611 Yeap, agree with the Python anti-pattern. I will make the changes and let's see what the the CI says, hopefully I made come mistake when testing and it will be able to resolve the None type.

rragundez · 2023-02-07T12:57:13Z

@kc611 done, please review. I think the CI is failing with an issue related to what I mentioned above, let me know what should I do to solve it or if I should revert back. Thanks.

rragundez · 2023-02-08T21:49:03Z

@kc611 By the way I just remembered that indeed is a bad practice to set default parameters that are mutable but even though numpy arrays are mutable, an array of size 0 is immutable because it has no elements to change, and numpy arrays are of fixed size, therefore the default value of an int array of size zero is immutable. Let me know if this makes sense and how you would like me to proceed.

kc611

Hi @rragundez , Thank you for making the changes. 😃 . I think the reason for failure is that Numba isn't able to infer I as None during compile time and ends up trying to lower the branches with argpartition boolean. A way around this would be to make argpartition a compile time constant instead of runtime, this would make sure branch pruner absolutely removes the branches with the given boolean so that they aren't lowered into IR.

I agree that an array of size 0 is immutable because it has no elements to change, but we can actually change it's attributes making it a 'mutable as a structure'. 🙂 However, the issue at hand is that this practice is generally not recommended and even more so in Numba because it might lead to weird and unexpected IR issues down the line, even if it seems to work just fine right now.

rragundez · 2023-02-11T15:21:25Z

@kc611 thanks for the feedback. I made the requested changes and the CI is passing now. Please review.

kc611

Thanks @rragundez, these changes LGTM!

rragundez · 2023-02-16T23:50:35Z

@kc611 Thanks for the approval. I cannot merge though, still says "Merging is blocked Merging can be performed automatically with 1 approving review."

rragundez · 2023-02-21T05:18:48Z

Hi @kc611, can you help me understand if this PR will be able to be merged? Thanks.

Add numpy argpartition function support

5192cdb

rragundez requested review from sklam and stuartarchibald as code owners January 28, 2023 06:58

kc611 added 3 - Ready for Review numpy Effort - medium Medium size effort needed labels Jan 30, 2023

kc611 approved these changes Jan 30, 2023

View reviewed changes

kc611 requested changes Jan 30, 2023

View reviewed changes

Comment thread numba/np/arraymath.py Outdated

Comment thread numba/np/arraymath.py Outdated

Comment thread numba/np/arraymath.py Outdated

guilhermeleobas assigned kc611 Jan 31, 2023

guilhermeleobas added 4 - Waiting on reviewer Waiting for reviewer to respond to author and removed 3 - Ready for Review labels Jan 31, 2023

kc611 added 4 - Waiting on author Waiting for author to respond to review and removed 4 - Waiting on reviewer Waiting for reviewer to respond to author labels Feb 7, 2023

Change default parameter to None

193ec50

Change logic based on default parameter

413f2d5

rragundez requested review from kc611 and removed request for sklam and stuartarchibald February 10, 2023 08:51

kc611 requested changes Feb 10, 2023

View reviewed changes

Comment thread numba/np/arraymath.py Outdated

Comment thread numba/np/arraymath.py Outdated

Comment thread numba/np/arraymath.py

Comment thread numba/np/arraymath.py

Comment thread numba/np/arraymath.py Outdated

Rodrigo Agundez added 2 commits February 10, 2023 19:27

Add boolean argpartition argument to partition factory

d5a66ed

Fix line length error from flake8

eefac1a

kc611 approved these changes Feb 16, 2023

View reviewed changes

stuartarchibald added 4 - Waiting on reviewer Waiting for reviewer to respond to author and removed 4 - Waiting on author Waiting for author to respond to review labels Feb 23, 2023

kc611 added 5 - Ready to merge Review and testing done, is ready to merge and removed 4 - Waiting on reviewer Waiting for reviewer to respond to author Effort - medium Medium size effort needed labels Feb 23, 2023

sklam added this to the Numba 0.57 RC milestone Feb 25, 2023

sklam approved these changes Feb 25, 2023

View reviewed changes

sklam merged commit de4cd6c into numba:main Feb 25, 2023

rragundez deleted the np-argpartition branch February 25, 2023 01:06

rragundez mentioned this pull request Feb 27, 2023

Please add np.argpartition implementation in nopython mode #2445

Closed

Conversation

rragundez commented Jan 28, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kc611 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kc611 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

rragundez commented Jan 30, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

rragundez commented Feb 6, 2023

Uh oh!

kc611 commented Feb 7, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

rragundez commented Feb 7, 2023

Uh oh!

rragundez commented Feb 7, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

rragundez commented Feb 8, 2023

Uh oh!

kc611 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

rragundez commented Feb 11, 2023

Uh oh!

kc611 left a comment

Choose a reason for hiding this comment

Uh oh!

rragundez commented Feb 16, 2023

Uh oh!

rragundez commented Feb 21, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

rragundez commented Jan 28, 2023 •

edited

Loading

kc611 left a comment •

edited

Loading

rragundez commented Jan 30, 2023 •

edited

Loading

kc611 commented Feb 7, 2023 •

edited

Loading

rragundez commented Feb 7, 2023 •

edited

Loading

kc611 left a comment •

edited

Loading