[BUG] mx.argmax / mx.argmin skip NaN while mx.max / mx.min propagate it

[ayaangazali]

Describe the bug

mx.argmax and mx.argmin skip NaN, while mx.max and mx.min propagate it. So on any array containing a NaN, a[mx.argmax(a)] and mx.max(a) disagree:

>>> a = mx.array([3.0, float("nan"), 1.0, 5.0])
>>> mx.max(a)
array(nan, dtype=float32)
>>> mx.argmax(a)
array(3, dtype=uint32)     # a[3] is 5.0, not nan

numpy and pytorch both return the index of the first NaN, which keeps argmax and max consistent with each other.

case	mlx argmax	numpy	torch	mlx argmin	numpy	torch
[3, nan, 1, 5]	3	1	1	2	1	1
[nan, 1, 2]	0	0	0	0	0	0
[1, 2, nan]	1	2	2	0	2	2
[inf, nan, -inf]	0	1	1	2	1	1
[-1, -2, nan]	0	2	2	1	2	2

This is the same shape as #4044, which was fixed for cummax / cummin, and #4072, which is still open for all_max / all_min. mx.cummax propagates NaN today and mx.max does too, so argmax and argmin are the remaining pair that does not.

Not size dependent, so it is not a vector versus scalar tail issue: with the NaN in the middle, lengths 3, 4, 8, 33, 1000 and 5000 all return the index of the largest non NaN. Same for float16 and bfloat16, and the same along an explicit axis.

To Reproduce

import mlx.core as mx

a = mx.array([3.0, float("nan"), 1.0, 5.0])
print(mx.max(a), mx.argmax(a), a[mx.argmax(a).item()])
# array(nan, dtype=float32) array(3, dtype=uint32) array(5, dtype=float32)

Expected behavior

a[mx.argmax(a)] == mx.max(a) whenever the array is non empty, including when it holds a NaN, which means returning the index of the first NaN. Same for argmin and min.

Desktop

• OS Version: macOS 15.5 (Darwin 25.5.0)
• Version: 0.32.1.dev20260815+9ab977b56, CPU only build

Additional context

The CPU side is mlx/backend/cpu/arg_reduce.cpp, where the comparison is a plain x > (*y), so a NaN never wins. Metal and CUDA use a pairwise tree reduction with a best.index > current.index tie break, so making the first NaN win there needs the NaN cases handled explicitly in the combine rather than just a comparison change.

I am happy to send a PR across all three backends if you want the numpy and pytorch behaviour. I asked rather than just sending one because it changes a defined result on every backend, and I have no Metal or CUDA machine here, so I would be leaning on CI to check the two kernel paths.

I am a freshman learning this codebase and I used Claude Code while investigating. Every number above came from a run on this machine.