Skip to content

ENH: add support for np.nanargmax and np.nanargmin - #10808

Merged
guilhermeleobas merged 7 commits into
numba:mainfrom
juanmicl:add-np-nanargmax-nanargmin
Sep 8, 2026
Merged

guilhermeleobas merged 7 commits into
numba:mainfrom
juanmicl:add-np-nanargmax-nanargmin

Conversation

@juanmicl

Copy link
Copy Markdown
Contributor

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.nanargmin have never been supported in nopython
mode; jitting them fails at typing time:

numba.core.errors.TypingError: Failed in nopython mode pipeline
(step: nopython frontend)
Use of unsupported NumPy function 'numpy.nanargmax' or unsupported use
of the function.

What changed

  • Adds @overload(np.nanargmax) and @overload(np.nanargmin) to
    numba/np/arraymath.py. Both share a small factory,
    build_array_nanarg_impl(name, operator), that returns a
    register_jitable implementation doing a single NaN-skipping pass over
    the 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.argmin implementations, since NaN is not representable
    there and NumPy treats such input like the plain functions. The axis
    argument is handled by reusing the same helper those functions use,
    build_argmax_or_argmin_with_axis_impl. The
    numpy_version < (1, 22) guard from the original sketch is no longer
    needed as main already requires 1.22.

  • Tests go into numba/tests/test_array_reductions.py as one combined
    parametrised test covering float and integer dtypes, axis=None plus
    positive/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):

$ ./runtests.py numba.tests.test_array_reductions.TestArrayReductions.test_nanargmax_nanargmin -m 1
FAILED (errors=1)
# numba.core.errors.TypingError: Use of unsupported NumPy function
# 'numpy.nanargmax' or unsupported use of the function.

Green after the implementation:

$ ./runtests.py numba.tests.test_array_reductions.TestArrayReductions.test_nanargmax_nanargmin -m 1
Ran 1 test in 30.073s
OK

Full touched-file suite:

$ ./runtests.py numba.tests.test_array_reductions -m 8
Ran 448 tests in 93.130s
OK (skipped=1)

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

  • The flatten impl keeps a single pass with a sentinel (max_idx = -1);
    the first non-NaN element is taken unconditionally via the max_idx < 0
    short-circuit.
  • Ties keep the first occurrence, matching np.argmax/np.argmin
    semantics.

@guilhermeleobas guilhermeleobas left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@guilhermeleobas guilhermeleobas added the 4 - Waiting on author Waiting for author to respond to review label Sep 4, 2026
@juanmicl

juanmicl commented Sep 7, 2026

Copy link
Copy Markdown
Contributor Author

Both items are addressed:

  • numpysupported.rst now lists np.nanargmax / np.nanargmin.
  • The test now covers non-contiguous and F-order input (the latter both flattened and with an explicit axis), 0-d arrays, and ±inf values, following the coverage of NumPy's own nanfunctions tests.

The ±inf cases exposed a real divergence in the original implementation: NumPy implements nanarg* by replacing NaN with ±inf and then running a plain argmax/argmin, so for example np.nanargmin([nan, inf]) returns 0, the NaN position wins the tie against the fill value. The implementation now mimics that replacement strategy. I verified the behaviour matches on NumPy 1.23, 2.0 and 2.5.

One limitation, noted in a comment in the test: 0-d input with an explicit axis is not supported, consistent with the existing np.argmax/np.argmin (they share the same axis helper).

@guilhermeleobas guilhermeleobas added 5 - Ready to merge Review and testing done, is ready to merge and removed 3 - Ready for Review 4 - Waiting on author Waiting for author to respond to review labels Sep 8, 2026

@guilhermeleobas guilhermeleobas left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @juanmicl

@guilhermeleobas
guilhermeleobas merged commit 5e70dbf into numba:main Sep 8, 2026
24 of 27 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

5 - Ready to merge Review and testing done, is ready to merge

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants