inline complex isnan

Author: Lyxot

mlx/backend/metal/kernels/complex.h

@@ -97,10 +97,6 @@ constexpr bool operator<(complex64_t a, complex64_t b) {
   return operator>(b, a);
 }
 
-constexpr bool isnan(complex64_t x) {
-  return isnan(x.real) || isnan(x.imag);
-}
-
 constexpr bool operator==(complex64_t a, complex64_t b) {
   return a.real == b.real && a.imag == b.imag;
 }

mlx/backend/metal/kernels/sort.h

@@ -40,10 +40,15 @@ template <typename T>
 struct LessThan {
   static constexpr constant T init = Init<T>::v;
   METAL_FUNC bool operator()(T a, T b) const {
-    if constexpr (
-        metal::is_floating_point_v<T> || metal::is_same_v<T, complex64_t>) {
-      bool an = isnan(a);
-      bool bn = isnan(b);
+    if constexpr (metal::is_floating_point_v<T>) {
+      bool an = metal::isnan(a);
+      bool bn = metal::isnan(b);
+      if (an | bn) {
+        return (!an) & bn;
+      }
+    } else if constexpr (metal::is_same_v<T, complex64_t>) {
+      bool an = metal::isnan(a.real) || metal::isnan(a.imag);
+      bool bn = metal::isnan(b.real) || metal::isnan(b.imag);
       if (an | bn) {
         return (!an) & bn;
       }