DOC: Add take_along_axis to the see also section in argmin, argmax etc. #14799
Conversation
…t and argpartition.
numpy/core/fromnumeric.py
Outdated
| @@ -797,6 +797,7 @@ def argpartition(a, kth, axis=-1, kind='introselect', order=None): | |||
| partition : Describes partition algorithms used. | |||
| ndarray.partition : Inplace partition. | |||
| argsort : Full indirect sort | |||
| take_along_axis : Take values from the input array by matching 1d index and data slices. | |||
There was a problem hiding this comment.
Let's see how another maintainer feels about this first, but perhaps:
| take_along_axis : Take values from the input array by matching 1d index and data slices. | |
| take_along_axis : Apply ``index_array`` from argpartition to an array as if by calling partition |
And same for the below,
take_along_axis : Apply ``index_array`` from argsort to an array as if by calling sort
take_along_axis : Apply ``np.expand_dims(index_array, axis)`` from argmin to an array as if by calling min
take_along_axis : Apply ``np.expand_dims(index_array, axis)`` from argmax to an array as if by calling max
There was a problem hiding this comment.
Probably those use cases are worth mentioning, but I wasn't sure if the "See also" section is a right place for that. e.g. argsort and argpartition refer to take_along_axis in "Returns". Perhaps it's better to just add an example.
There was a problem hiding this comment.
An example for each of these showing how to use take_along_axis would be great, if you're willing to spend the time. If not, we can merge this as is, and leave the issue open for someone else to add the examples.
There was a problem hiding this comment.
"See Also" seems appropriate here.
numpy/core/fromnumeric.py
Outdated
| @@ -1223,6 +1239,13 @@ def argmin(a, axis=None, out=None): | |||
| >>> np.argmin(b) # Only the first occurrence is returned. | |||
| 0 | |||
|
|
|||
| >>> x = np.array([[4,2,3],[1,0,3]]) | |||
| >>> np.take_along_axis(x, np.expand_dims(np.argmin(x, axis=-1), axis=-1), axis=-1) | |||
There was a problem hiding this comment.
Possibly worth splitting these examples in two, either as:
| >>> np.take_along_axis(x, np.expand_dims(np.argmin(x, axis=-1), axis=-1), axis=-1) | |
| >>> index_array = np.argmin(x, axis=-1) | |
| >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1) |
or
| >>> np.take_along_axis(x, np.expand_dims(np.argmin(x, axis=-1), axis=-1), axis=-1) | |
| >>> index_array = np.expand_dims(np.argmin(x, axis=-1), axis=-1) # emulate keepdims=True | |
| >>> np.take_along_axis(x, index_array, axis=-1) |
Co-Authored-By: Eric Wieser <wieser.eric@gmail.com>
Co-Authored-By: Eric Wieser <wieser.eric@gmail.com>
Co-Authored-By: Eric Wieser <wieser.eric@gmail.com>
Co-Authored-By: Eric Wieser <wieser.eric@gmail.com>
Co-Authored-By: Eric Wieser <wieser.eric@gmail.com>
Co-Authored-By: Eric Wieser <wieser.eric@gmail.com>
Co-Authored-By: Eric Wieser <wieser.eric@gmail.com>
numpy/core/fromnumeric.py
Outdated
| >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1) # same as np.min(x, axis=-1, keepdims=True) | ||
| array([[2], | ||
| [0]]) | ||
| >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1).squeeze(axis=-1) # same as np.max(x, axis=-1) |
There was a problem hiding this comment.
Nit: pep8 wants two spaces before #, same on other lines.
There was a problem hiding this comment.
Thanks, I was looking for that bug
| [3]]) | ||
| >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1).squeeze(axis=-1) # same as np.max(x, axis=-1) | ||
| array([4, 3]) | ||
|
|
There was a problem hiding this comment.
These lines are quite long. Maybe move the comment to stand alone in the line above the code
|
A small nit about line length, but otherwise this looks good |
|
Thanks @mproszewska |
This closes #6078.