Remove tuple unpacking syntax to comply with earlier python versions. Add cuda skip to nanpropagation tests, fix cuda implementation in a separate PR.

mlx/backend/cpu/reduce.cpp

@@ -325,7 +325,15 @@ struct MaxReduce {
   };
 
   template <int N, typename T>
-  T operator()(simd::Simd<T, N> x) {
+  std::enable_if_t<std::is_integral_v<T>, T> operator()(simd::Simd<T, N> x) {
+    return simd::max(x);
+  };
+
+  template <int N, typename T>
+  std::enable_if_t<!std::is_integral_v<T>, T> operator()(simd::Simd<T, N> x) {
+    if (simd::any(x != x)) {
+      return static_cast<T>(NAN);
+    }
     return simd::max(x);
   };
 };
@@ -527,10 +535,10 @@ void Reduce::eval_cpu(const std::vector<array>& inputs, array& out) {
             reduce_dispatch_min_max<uint64_t>(in, out, reduce_type_, axes_);
             break;
           case int8:
-            reduce_dispatch_min_max<uint8_t>(in, out, reduce_type_, axes_);
+            reduce_dispatch_min_max<int8_t>(in, out, reduce_type_, axes_);
             break;
           case int16:
-            reduce_dispatch_min_max<uint16_t>(in, out, reduce_type_, axes_);
+            reduce_dispatch_min_max<int16_t>(in, out, reduce_type_, axes_);
             break;
           case int32:
             reduce_dispatch_min_max<int32_t>(in, out, reduce_type_, axes_);

python/tests/cuda_skip.py

@@ -3,6 +3,8 @@ cuda_skip = {
     "TestLayers.test_quantized_embedding",
     "TestOps.test_dynamic_slicing",
     "TestReduce.test_dtypes",
+    "TestReduce.test_nanpropagation",
+    "TestReduce.test_nanpropagation_complex64",
     # Block masked matmul NYI
     "TestBlas.test_block_masked_matmul",
     # Gather matmul NYI

python/tests/test_reduce.py

@@ -167,10 +167,10 @@ class TestReduce(mlx_tests.MLXTestCase):
 
         for dtype in dtypes:
             with self.subTest(dtype=dtype):
-                x = (mx.random.normal((4, 4))).astype(getattr(mx, dtype))
+                x = (mx.random.normal((4, 4)) * 10).astype(getattr(mx, dtype))
                 indices = mx.random.randint(0, 4, shape=(6,)).reshape(3, 2)
                 for idx in indices:
-                    x[*idx] = mx.nan
+                    x[idx[0], idx[1]] = mx.nan
                 x_np = np.array(x)
 
                 for op in ["max"]: