fix: propagate NaN in arg reductions

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"
@@ -11,6 +12,14 @@ namespace mlx::core {
 
 namespace {
 
+template <typename T>
+bool is_nan(T x) {
+  if constexpr (is_floating_point_v<T>) {
+    return std::isnan(x);
+  }
+  return false;
+}
+
 template <typename InT, typename OpT>
 void arg_reduce(const array& in, array& out, const OpT& op, int axis) {
   auto axis_size = in.shape()[axis];
@@ -41,7 +50,7 @@ void arg_reduce_dispatch(
   switch (rtype) {
     case ArgReduce::ArgMin: {
       auto op = [](auto ind_x, auto x, auto ind_y, auto y) {
-        if (x < (*y)) {
+        if ((!is_nan(*y) && is_nan(x)) || x < (*y)) {
           (*y) = x;
           (*ind_y) = ind_x;
         }
@@ -51,7 +60,7 @@ void arg_reduce_dispatch(
     }
     case ArgReduce::ArgMax: {
       auto op = [](auto ind_x, auto x, auto ind_y, auto y) {
-        if (x > (*y)) {
+        if ((!is_nan(*y) && is_nan(x)) || x > (*y)) {
           (*y) = x;
           (*ind_y) = ind_x;
         }

mlx/backend/cuda/arg_reduce.cu

@@ -26,6 +26,14 @@ struct IndexValPair {
   T val;
 };
 
+template <typename T>
+__device__ bool is_nan(T x) {
+  if constexpr (is_floating_v<T>) {
+    return cuda::std::isnan(x);
+  }
+  return false;
+}
+
 template <typename T>
 struct ArgMin {
   constexpr __device__ T init() {
@@ -35,8 +43,11 @@ struct ArgMin {
   __device__ IndexValPair<T> operator()(
       const IndexValPair<T>& best,
       const IndexValPair<T>& current) {
-    if (best.val > current.val ||
-        (best.val == current.val && best.index > current.index)) {
+    if ((is_nan(current.val) &&
+         (!is_nan(best.val) || best.index > current.index)) ||
+        (!is_nan(best.val) &&
+         (best.val > current.val ||
+          (best.val == current.val && best.index > current.index)))) {
       return current;
     } else {
       return best;
@@ -50,7 +61,7 @@ struct ArgMin {
       uint32_t offset) {
 #pragma unroll
     for (int i = 0; i < N; i++) {
-      if (vals[i] < best.val) {
+      if ((!is_nan(best.val) && is_nan(vals[i])) || vals[i] < best.val) {
         best.val = vals[i];
         best.index = offset + i;
       }
@@ -68,8 +79,11 @@ struct ArgMax {
   __device__ IndexValPair<T> operator()(
       const IndexValPair<T>& best,
       const IndexValPair<T>& current) {
-    if (best.val < current.val ||
-        (best.val == current.val && best.index > current.index)) {
+    if ((is_nan(current.val) &&
+         (!is_nan(best.val) || best.index > current.index)) ||
+        (!is_nan(best.val) &&
+         (best.val < current.val ||
+          (best.val == current.val && best.index > current.index)))) {
       return current;
     } else {
       return best;
@@ -83,7 +97,7 @@ struct ArgMax {
       uint32_t offset) {
 #pragma unroll
     for (int i = 0; i < N; i++) {
-      if (vals[i] > best.val) {
+      if ((!is_nan(best.val) && is_nan(vals[i])) || vals[i] > best.val) {
         best.val = vals[i];
         best.index = offset + i;
       }

mlx/backend/metal/kernels/arg_reduce.metal

@@ -12,13 +12,24 @@ struct IndexValPair {
   U val;
 };
 
+template <typename U>
+bool is_nan(U x) {
+  if constexpr (metal::is_floating_point_v<U>) {
+    return isnan(x);
+  }
+  return false;
+}
+
 template <typename U>
 struct ArgMin {
   static constexpr constant U init = Limits<U>::max;
 
   IndexValPair<U> reduce(IndexValPair<U> best, IndexValPair<U> current) thread {
-    if (best.val > current.val ||
-        (best.val == current.val && best.index > current.index)) {
+    if ((is_nan(current.val) &&
+         (!is_nan(best.val) || best.index > current.index)) ||
+        (!is_nan(best.val) &&
+         (best.val > current.val ||
+          (best.val == current.val && best.index > current.index)))) {
       return current;
     } else {
       return best;
@@ -29,7 +40,7 @@ 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 ((!is_nan(best.val) && is_nan(vals[i])) || vals[i] < best.val) {
         best.val = vals[i];
         best.index = offset + i;
       }
@@ -43,8 +54,11 @@ struct ArgMax {
   static constexpr constant U init = Limits<U>::min;
 
   IndexValPair<U> reduce(IndexValPair<U> best, IndexValPair<U> current) thread {
-    if (best.val < current.val ||
-        (best.val == current.val && best.index > current.index)) {
+    if ((is_nan(current.val) &&
+         (!is_nan(best.val) || best.index > current.index)) ||
+        (!is_nan(best.val) &&
+         (best.val < current.val ||
+          (best.val == current.val && best.index > current.index)))) {
       return current;
     } else {
       return best;
@@ -55,7 +69,7 @@ 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 ((!is_nan(best.val) && is_nan(vals[i])) || vals[i] > best.val) {
         best.val = vals[i];
         best.index = offset + i;
       }

tests/arg_reduce_tests.cpp

@@ -1,5 +1,7 @@
 // Copyright © 2023 Apple Inc.
 
+#include <limits>
+
 #include "doctest/doctest.h"
 
 #include "mlx/mlx.h"
@@ -187,6 +189,31 @@ TEST_CASE("test arg reduce edge cases") {
   CHECK_THROWS(argmax(array({})));
 }
 
+TEST_CASE("test arg reduce NaN") {
+  auto nan = std::numeric_limits<float>::quiet_NaN();
+  auto x = array({3.0f, nan, 1.0f, 5.0f, nan, -1.0f}, {2, 3});
+
+  test_arg_reduce_small(Device::cpu, x, ArgReduce::ArgMin, {2}, 1, {1, 1});
+  test_arg_reduce_small(Device::cpu, x, ArgReduce::ArgMax, {2}, 1, {1, 1});
+
+  if (!metal::is_available()) {
+    INFO("Skipping arg reduction gpu tests");
+    return;
+  }
+
+  test_arg_reduce_small(Device::gpu, x, ArgReduce::ArgMin, {2}, 1, {1, 1});
+  test_arg_reduce_small(Device::gpu, x, ArgReduce::ArgMax, {2}, 1, {1, 1});
+
+  std::vector<float> wide(1024, 0.0f);
+  wide[17] = nan;
+  wide[529] = nan;
+  x = array(wide.data(), {1024});
+  test_arg_reduce_small(Device::cpu, x, ArgReduce::ArgMin, {}, 0, {17});
+  test_arg_reduce_small(Device::cpu, x, ArgReduce::ArgMax, {}, 0, {17});
+  test_arg_reduce_small(Device::gpu, x, ArgReduce::ArgMin, {}, 0, {17});
+  test_arg_reduce_small(Device::gpu, x, ArgReduce::ArgMax, {}, 0, {17});
+}
+
 TEST_CASE("test arg reduce irregular strides") {
   auto x = array(
       {0, 2, 1, 7, 5, -5, 0, 2, 1, 7, 5, -5,