[Metal] Support sorting complex numbers (#3314)

mlx/backend/metal/kernels/sort.h

@@ -29,14 +29,26 @@ struct Init<T, metal::enable_if_t<metal::is_floating_point_v<T>>> {
   static constexpr constant T v = metal::numeric_limits<T>::quiet_NaN();
 };
 
+template <>
+struct Init<complex64_t> {
+  static constexpr constant complex64_t v = complex64_t(
+      metal::numeric_limits<float>::quiet_NaN(),
+      metal::numeric_limits<float>::quiet_NaN());
+};
+
 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;
       }

mlx/backend/metal/kernels/sort.metal

@@ -50,6 +50,7 @@ instantiate_block_sort_bn(bfloat16, bfloat16_t)
 
 instantiate_block_sort_long(uint64, uint64_t)
 instantiate_block_sort_long(int64, int64_t)
+instantiate_block_sort_long(complex64, complex64_t)
 
 #define instantiate_multi_block_sort(                                      \
     vtname, vtype, itname, itype, arg_sort, bn, tn)                        \
@@ -77,4 +78,5 @@ instantiate_multi_block_sort_base(bfloat16, bfloat16_t)
   instantiate_multi_block_sort(vtname, vtype, uint32, uint32_t, true, 256, 4)
 
 instantiate_multi_block_sort_long(uint64, uint64_t)
-instantiate_multi_block_sort_long(int64, int64_t) // clang-format on
+instantiate_multi_block_sort_long(int64, int64_t)
+instantiate_multi_block_sort_long(complex64, complex64_t) // clang-format on

python/tests/test_ops.py

@@ -2192,11 +2192,8 @@ class TestOps(mlx_tests.MLXTestCase):
 
     def test_sort(self):
         shape = (6, 4, 10)
-        dtypes = ["int32", "float32"]
-        if not mx.metal.is_available():
-            dtypes.append("complex64")
         tests = product(
-            dtypes,  # type
+            ("int32", "float32", "complex64"),  # type
             (None, 0, 1, 2),  # axis
             (True, False),  # strided
         )
@@ -3326,10 +3323,9 @@ class TestOps(mlx_tests.MLXTestCase):
                 expected = mx.array([0.0, 2.0, 3.0, mx.nan], dtype=dtype)
                 self.assertTrue(mx.array_equal(mx.sort(x), expected, equal_nan=True))
 
-        if not mx.metal.is_available():
-            x = mx.array([3.0 + 1j, mx.nan + 2j, 2.0 + 1j, 0.0 + 1j])
-            expected = mx.array([0.0 + 1j, 2.0 + 1j, 3.0 + 1j, mx.nan + 2j])
-            self.assertTrue(mx.array_equal(mx.sort(x), expected, equal_nan=True))
+        x = mx.array([3.0 + 1j, mx.nan + 2j, 2.0 + 1j, 0.0 + 1j])
+        expected = mx.array([0.0 + 1j, 2.0 + 1j, 3.0 + 1j, mx.nan + 2j])
+        self.assertTrue(mx.array_equal(mx.sort(x), expected, equal_nan=True))
 
     def test_argsort_nan(self):
         for dtype in [mx.float32, mx.float16, mx.bfloat16]: