Repository navigation
ENH: add support for np.nanargmax and np.nanargmin - #10808
Conversation
guilhermeleobas
left a comment
There was a problem hiding this comment.
Code looks good to me. Can you update numpysupported.rst to include np.nanargmax and np.nanargmin?
Also, I asked claude to test your PR with the NumPy tests and I think it is worth to add the following cases:
- Non-contiguous and F-order input
- 0-d arrays
- arrays with +/- inf
You can find the tests in https://github.com/numpy/numpy/blob/3d065710eee404a34788deb8f05f977a88879b20/numpy/lib/tests/test_nanfunctions.py#L269
|
Both items are addressed:
The ±inf cases exposed a real divergence in the original implementation: NumPy implements One limitation, noted in a comment in the test: 0-d input with an explicit axis is not supported, consistent with the existing |
guilhermeleobas
left a comment
There was a problem hiding this comment.
Thanks @juanmicl
Part of #4074. Revives #8688 (original implementation by @Tobi995, closed
stale), reworked following the refactor suggested by @guilhermeleobas in
that review: a single shared implementation factory plus one combined
test, instead of two near-duplicated functions and tests.
Problem
np.nanargmax/np.nanargminhave never been supported innopythonmode; jitting them fails at typing time:
What changed
Adds
@overload(np.nanargmax)and@overload(np.nanargmin)tonumba/np/arraymath.py. Both share a small factory,build_array_nanarg_impl(name, operator), that returns aregister_jitableimplementation doing a single NaN-skipping pass overthe flattened array and raising NumPy's exact error messages for empty
and all-NaN input.
Integer and boolean dtypes simply delegate to the existing
np.argmax/np.argminimplementations, since NaN is not representablethere and NumPy treats such input like the plain functions. The
axisargument is handled by reusing the same helper those functions use,
build_argmax_or_argmin_with_axis_impl. Thenumpy_version < (1, 22)guard from the original sketch is no longerneeded as
mainalready requires 1.22.Tests go into
numba/tests/test_array_reductions.pyas one combinedparametrised test covering float and integer dtypes,
axis=Nonepluspositive/negative axes on 2-D and 4-D input, and the ValueError paths.
Complex and datetime64/timedelta64 dtypes are still unsupported, happy
to follow up with those if wanted. Also adds a release note under
docs/upcoming_changes/.Verification
TDD, red first (against the untouched tree):
Green after the implementation:
Full touched-file suite:
flake8 numba/np/arraymath.py numba/tests/test_array_reductions.py: clean.Environment: Python 3.12.8, numpy 2.5.2, llvmlite 0.50.0.dev2 (numba dev
index), numba main @ 3598bdd.
Notes for the reviewer
max_idx = -1);the first non-NaN element is taken unconditionally via the
max_idx < 0short-circuit.
np.argmax/np.argminsemantics.