ENH: Adding keepdims to np.argmin,np.argmax#19211
Conversation
rgommers
left a comment
There was a problem hiding this comment.
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. |
BvB93
left a comment
There was a problem hiding this comment.
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. |
seberg
left a comment
There was a problem hiding this comment.
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