From c405d0c17cc6507f883b2e412fcf67b58439ad57 Mon Sep 17 00:00:00 2001 From: Erwin Zhang Date: Sun, 16 Aug 2026 01:01:02 -0400 Subject: [PATCH] Return the index of the first NaN from argmax and argmin Every comparison against a NaN is false, so the bare < and > in the arg reduce operators step over NaNs entirely. argmax and argmin return the index of the largest or smallest real value while max and min return NaN for the same input, so a[mx.argmax(a)] and mx.max(a) disagree on any array containing one. numpy and torch both return the index of the first NaN. MLX already propagates NaN in max, min, cummax and cummin, the last two since #4044, so argmax and argmin were the remaining pair that did not agree with the rest. A NaN now wins the comparison, and once one is held it is not replaced, so the index is the first NaN rather than whichever the reduction happened to associate last. The float check is a compile time guard, so the integer instantiations are unchanged and still compile: cpu uses is_floating_point_v from half_types.h, which covers float16 and bfloat16, metal and cuda use a small arg_is_nan helper around the same idea. reduce_many needed the same treatment as reduce, or the vectorised path would skip NaNs for longer arrays while the scalar path did not. Verified on an M5 Max, cpu and metal, against the five cases in the report plus float16 and bfloat16, lengths 8, 33, 1000 and 5000 to cover the vector path, an array with two NaNs to check the first one wins, and integer input. Without a NaN nothing changes, including the lowest index tie break. test_reduce, test_ops and test_autograd pass. The cuda path is the same change but is not built here. --- mlx/backend/cpu/arg_reduce.cpp | 27 +++++++++++++++ mlx/backend/cuda/arg_reduce.cu | 39 ++++++++++++++++++++-- mlx/backend/metal/kernels/arg_reduce.metal | 39 ++++++++++++++++++++-- python/tests/test_reduce.py | 39 ++++++++++++++++++++++ 4 files changed, 140 insertions(+), 4 deletions(-) diff --git a/mlx/backend/cpu/arg_reduce.cpp b/mlx/backend/cpu/arg_reduce.cpp index d93f74cf0a..290e5194dd 100644 --- a/mlx/backend/cpu/arg_reduce.cpp +++ b/mlx/backend/cpu/arg_reduce.cpp @@ -1,6 +1,7 @@ // Copyright © 2023 Apple Inc. #include +#include #include "mlx/backend/common/utils.h" #include "mlx/backend/cpu/encoder.h" @@ -41,6 +42,21 @@ void arg_reduce_dispatch( switch (rtype) { case ArgReduce::ArgMin: { auto op = [](auto ind_x, auto x, auto ind_y, auto y) { + // Every comparison against a NaN is false, so a bare `x < *y` steps + // over NaNs entirely and returns the index of the smallest real value + // while `mx.min` returns NaN for the same input. Propagate instead, + // and keep the first NaN so the index is stable, which is what numpy + // and torch return. + if constexpr (is_floating_point_v) { + if (std::isnan(*y)) { + return; + } + if (std::isnan(x)) { + (*y) = x; + (*ind_y) = ind_x; + return; + } + } if (x < (*y)) { (*y) = x; (*ind_y) = ind_x; @@ -51,6 +67,17 @@ void arg_reduce_dispatch( } case ArgReduce::ArgMax: { auto op = [](auto ind_x, auto x, auto ind_y, auto y) { + // See ArgMin above: NaN wins, and the first one keeps the index. + if constexpr (is_floating_point_v) { + if (std::isnan(*y)) { + return; + } + if (std::isnan(x)) { + (*y) = x; + (*ind_y) = ind_x; + return; + } + } if (x > (*y)) { (*y) = x; (*ind_y) = ind_x; diff --git a/mlx/backend/cuda/arg_reduce.cu b/mlx/backend/cuda/arg_reduce.cu index 9609b85972..a65476c940 100644 --- a/mlx/backend/cuda/arg_reduce.cu +++ b/mlx/backend/cuda/arg_reduce.cu @@ -26,6 +26,19 @@ struct IndexValPair { T val; }; +// Every comparison against a NaN is false, so a plain `<` or `>` walks past +// NaNs and argmin/argmax return the index of the smallest or largest real +// value while min/max return NaN for the same input. Prefer the NaN, and keep +// the first one so the index is stable, matching numpy and torch. +template +__device__ inline bool arg_is_nan(T x) { + if constexpr (cuda::std::is_floating_point_v || is_floating_point_v) { + return cuda::std::isnan(x); + } else { + return false; + } +} + template struct ArgMin { constexpr __device__ T init() { @@ -35,6 +48,14 @@ struct ArgMin { __device__ IndexValPair operator()( const IndexValPair& best, const IndexValPair& current) { + bool bn = arg_is_nan(best.val); + bool cn = arg_is_nan(current.val); + if (bn || cn) { + if (bn && cn) { + return best.index <= current.index ? best : current; + } + return bn ? best : current; + } if (best.val > current.val || (best.val == current.val && best.index > current.index)) { return current; @@ -50,7 +71,10 @@ struct ArgMin { uint32_t offset) { #pragma unroll for (int i = 0; i < N; i++) { - if (vals[i] < best.val) { + if (arg_is_nan(best.val)) { + break; + } + if (arg_is_nan(vals[i]) || vals[i] < best.val) { best.val = vals[i]; best.index = offset + i; } @@ -68,6 +92,14 @@ struct ArgMax { __device__ IndexValPair operator()( const IndexValPair& best, const IndexValPair& current) { + bool bn = arg_is_nan(best.val); + bool cn = arg_is_nan(current.val); + if (bn || cn) { + if (bn && cn) { + return best.index <= current.index ? best : current; + } + return bn ? best : current; + } if (best.val < current.val || (best.val == current.val && best.index > current.index)) { return current; @@ -83,7 +115,10 @@ struct ArgMax { uint32_t offset) { #pragma unroll for (int i = 0; i < N; i++) { - if (vals[i] > best.val) { + if (arg_is_nan(best.val)) { + break; + } + if (arg_is_nan(vals[i]) || vals[i] > best.val) { best.val = vals[i]; best.index = offset + i; } diff --git a/mlx/backend/metal/kernels/arg_reduce.metal b/mlx/backend/metal/kernels/arg_reduce.metal index 6ce20555f7..36e9bc3ff3 100644 --- a/mlx/backend/metal/kernels/arg_reduce.metal +++ b/mlx/backend/metal/kernels/arg_reduce.metal @@ -6,6 +6,19 @@ using namespace metal; +// Every comparison against a NaN is false, so a plain `<` or `>` walks past +// NaNs and argmin/argmax return the index of the smallest or largest real +// value while min/max return NaN for the same input. Prefer the NaN, and keep +// the first one so the index is stable, matching numpy and torch. +template +METAL_FUNC bool arg_is_nan(U x) { + if constexpr (metal::is_floating_point_v) { + return metal::isnan(x); + } else { + return false; + } +} + template struct IndexValPair { uint32_t index; @@ -17,6 +30,14 @@ struct ArgMin { static constexpr constant U init = Limits::max; IndexValPair reduce(IndexValPair best, IndexValPair current) thread { + bool bn = arg_is_nan(best.val); + bool cn = arg_is_nan(current.val); + if (bn || cn) { + if (bn && cn) { + return best.index <= current.index ? best : current; + } + return bn ? best : current; + } if (best.val > current.val || (best.val == current.val && best.index > current.index)) { return current; @@ -29,7 +50,10 @@ struct ArgMin { IndexValPair reduce_many(IndexValPair best, thread U* vals, uint32_t offset) thread { for (int i = 0; i < N; i++) { - if (vals[i] < best.val) { + if (arg_is_nan(best.val)) { + break; + } + if (arg_is_nan(vals[i]) || vals[i] < best.val) { best.val = vals[i]; best.index = offset + i; } @@ -43,6 +67,14 @@ struct ArgMax { static constexpr constant U init = Limits::min; IndexValPair reduce(IndexValPair best, IndexValPair current) thread { + bool bn = arg_is_nan(best.val); + bool cn = arg_is_nan(current.val); + if (bn || cn) { + if (bn && cn) { + return best.index <= current.index ? best : current; + } + return bn ? best : current; + } if (best.val < current.val || (best.val == current.val && best.index > current.index)) { return current; @@ -55,7 +87,10 @@ struct ArgMax { IndexValPair reduce_many(IndexValPair best, thread U* vals, uint32_t offset) thread { for (int i = 0; i < N; i++) { - if (vals[i] > best.val) { + if (arg_is_nan(best.val)) { + break; + } + if (arg_is_nan(vals[i]) || vals[i] > best.val) { best.val = vals[i]; best.index = offset + i; } diff --git a/python/tests/test_reduce.py b/python/tests/test_reduce.py index 6ac8fc1504..4d9a0c4e5a 100644 --- a/python/tests/test_reduce.py +++ b/python/tests/test_reduce.py @@ -273,6 +273,45 @@ def test_and_or_negative_zero(self): getattr(np, op)(x_np, axis=1).tolist(), ) + def test_arg_reduce_nan(self): + # argmax/argmin used to walk past NaN and return the index of the + # largest or smallest real value, while max/min return NaN for the + # same input. numpy and torch both return the index of the first NaN. + cases = [ + ([3.0, float("nan"), 1.0, 5.0], 1), + ([float("nan"), 1.0, 2.0], 0), + ([1.0, 2.0, float("nan")], 2), + ([float("inf"), float("nan"), float("-inf")], 1), + ([-1.0, -2.0, float("nan")], 2), + ] + for vals, want in cases: + a = mx.array(vals) + self.assertEqual(mx.argmax(a).item(), want) + self.assertEqual(mx.argmin(a).item(), want) + + # Not a short-array or tail effect, and the same for the half types. + for n in [8, 33, 1000, 5000]: + vals = [float(i) for i in range(n)] + vals[5] = float("nan") + a = mx.array(vals) + self.assertEqual(mx.argmax(a).item(), 5) + self.assertEqual(mx.argmin(a).item(), 5) + for dtype in [mx.float32, mx.float16, mx.bfloat16]: + a = mx.array([3.0, float("nan"), 1.0, 5.0], dtype=dtype) + self.assertEqual(mx.argmax(a).item(), 1) + self.assertEqual(mx.argmin(a).item(), 1) + + # The first NaN wins, so the index does not depend on how the + # reduction happens to associate. + a = mx.array([1.0, float("nan"), 3.0, float("nan")]) + self.assertEqual(mx.argmax(a).item(), 1) + self.assertEqual(mx.argmin(a).item(), 1) + + # Without a NaN nothing changes, including the lowest-index tie break. + a = mx.array([1.0, 5.0, 5.0, 1.0]) + self.assertEqual(mx.argmax(a).item(), 1) + self.assertEqual(mx.argmin(a).item(), 0) + if __name__ == "__main__": mlx_tests.MLXTestRunner(failfast=True)