From c7af3016eb8846f7ed4132aedcdf728a331845a8 Mon Sep 17 00:00:00 2001
From: Joona Havukainen <joona.j.havukainen@gmail.com>
Date: Mon, 7 Jul 2025 18:24:30 -0700
Subject: [PATCH] Only check nans on non-integral types in simd_reduce_impl.

---
 mlx/backend/metal/kernels/reduction/ops.h | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/mlx/backend/metal/kernels/reduction/ops.h b/mlx/backend/metal/kernels/reduction/ops.h
index ee6fc5e5f..57ddffef8 100644
--- a/mlx/backend/metal/kernels/reduction/ops.h
+++ b/mlx/backend/metal/kernels/reduction/ops.h
@@ -186,7 +186,12 @@ struct Max {
   DEFINE_SIMD_REDUCE()
 
   template <typename T>
-  T simd_reduce_impl(T val) {
+  metal::enable_if_t<metal::is_integral_v<T>, T> simd_reduce_impl(T val) {
+    return simd_max(val);
+  }
+
+  template <typename T>
+  metal::enable_if_t<!metal::is_integral_v<T>, T> simd_reduce_impl(T val) {
     if (simd_any(val != val)) {
       return static_cast<T>(NAN);
     }