-
-
Notifications
You must be signed in to change notification settings - Fork 10.3k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
ENH: Adding keepdims to np.argmin,np.argmax #19211
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks @czgdp1807.
This indeed looks deceptively simple - if it works, then it'd be fine with me to keep this PR in Python. Support in the C API could be added as a follow-up step.
The axis=None, keepdims=True
case doesn't seem tested though, and it looks like that'll need some special handling.
Agreed.
Ah! Yes. Missed it. Thanks for pointing out. I am on it. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can you update the argmax
/argmin
annotation as well (see fromnumeric.pyi)?
It should be fairly straightforward by the looks of things.
@overload
def argmax(
a: ArrayLike,
axis: None = ...,
out: Optional[ndarray] = ...,
+ keepdims: Literal[False] = ...,
) -> intp: ...
@overload
def argmax(
a: ArrayLike,
axis: Optional[int] = ...,
out: Optional[ndarray] = ...,
+ keepdims: bool = ...,
) -> Any: ...
Please let me know if there are any cases where the current approach fails or is inefficient. |
Ah, other issue: this introduces a mismatch between |
Yeah. I was thinking so, while making the change. Noticed that, |
I have made some updates to the interface of |
Did you mean to touch |
I thought that, Do we need to use, |
Ah no, if it was intentional then it's all good! Updating |
Yes, it looks like that's necessary. It looks odd; the reason for it is to support (partial) subclasses. Say |
@BvB93 I have added a release notes file to the PR. Are there any improvements to be made to that? |
Edit - Apologies for resolving the conversations here. Actually, I do it just keep track of stuff which is yet to be addressed. Let me know if I should un-resolve them. |
I suspect you might be confusing keyword-only ( In any case, compared to the annotations prior to 9260d40 only a single change will have to be made: @overload
def argmin(
a: ArrayLike,
axis: None = ...,
out: Optional[ndarray] = ...,
out: Optional[ndarray] = ...,
+ *,
+ keepdims: Literal[False] = ...,
) -> intp: ...
@overload
def argmin(
a: ArrayLike,
axis: Optional[int] = ...,
out: Optional[ndarray] = ...,
+ *,
+ keepdims: bool = ...,
) -> Any: ... |
I have restored the annotations. Let me know if anything else is required. Thanks. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
OK, a few nitpicks, but most of them are "take it or leave it". The test looks very thorough. Unfortunately it is a bit unwieldy, but I am OK with leaving it as it is. Thanks for bearing with me trying to find a minimal solution!
One awesome thing (happy with followup or skipping), would be to unify the two functions completely. They seem 100% copy-paste of each other with 2 variations (the method, and the "argmax" vs. "argmin" in error messages). And with PyErr_Format
these days, both of those should be straight forward to unify.
It might be good to ping the mailing list briefly, also about the C-API addition (which I honestly don't care about personally). Those should maybe also have a very brief (just a bullet point is enough) additional release note under the C-API section.
Anyway, this looks good to me, thanks Gagandeep!
Co-authored-by: Sebastian Berg <sebastian@sipsolutions.net>
I am thinking to do bit of refactoring (in C and test files) after this PR gets in. Let me know if that works.
I have added a release notes entry for C-API too. Let me know if it is okay.
In my opinion the solution is really elegant, simple and abstract. Thanks for guiding me to finally arrive at it. :-) |
I have removed the C-API label as well. Thanks. |
Co-authored-by: Eric Wieser <wieser.eric@gmail.com>
Co-authored-by: Eric Wieser <wieser.eric@gmail.com>
OK, lets put this in. Thanks Gagandeep! @czgdp1807, I think you did not apply all style fixes in both places always. I assume you will followup with a PR to merge the two functions into a single one that will clean up the styles where it is nicer? |
Yes sure. On the way. I will try to factor out the common parts in |
Closes #8710
Comments
The approach is too straightforward/trivial. Please let me know if the changes should be made at C-level, somewhere in
PyArray_ArgMax
.ping @eric-wieser @shoyer
cc @rgommers