Return the index of the first NaN from argmax and argmin

mlx/backend/cpu/arg_reduce.cpp

@@ -1,6 +1,7 @@
 // Copyright © 2023 Apple Inc.
 
 #include <cassert>
+#include <cmath>
 
 #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<InT>) {
+          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<InT>) {
+          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;

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 <typename T>
+__device__ inline bool arg_is_nan(T x) {
+  if constexpr (cuda::std::is_floating_point_v<T> || is_floating_point_v<T>) {
+    return cuda::std::isnan(x);
+  } else {
+    return false;
+  }
+}
+
 template <typename T>
 struct ArgMin {
   constexpr __device__ T init() {
@@ -35,6 +48,14 @@ struct ArgMin {
   __device__ IndexValPair<T> operator()(
       const IndexValPair<T>& best,
       const IndexValPair<T>& 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<T> operator()(
       const IndexValPair<T>& best,
       const IndexValPair<T>& 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;
       }

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 <typename U>
+METAL_FUNC bool arg_is_nan(U x) {
+  if constexpr (metal::is_floating_point_v<U>) {
+    return metal::isnan(x);
+  } else {
+    return false;
+  }
+}
+
 template <typename U>
 struct IndexValPair {
   uint32_t index;
@@ -17,6 +30,14 @@ struct ArgMin {
   static constexpr constant U init = Limits<U>::max;
 
   IndexValPair<U> reduce(IndexValPair<U> best, IndexValPair<U> 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<U>
   reduce_many(IndexValPair<U> 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<U>::min;
 
   IndexValPair<U> reduce(IndexValPair<U> best, IndexValPair<U> 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<U>
   reduce_many(IndexValPair<U> 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;
       }

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)