[CUDA] support sorting complex numbers (#3286)

python/tests/test_ops.py

@@ -2192,8 +2192,11 @@ 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(
-            ("int32", "float32"),  # type
+            dtypes,  # type
             (None, 0, 1, 2),  # axis
             (True, False),  # strided
         )
@@ -2201,7 +2204,13 @@ class TestOps(mlx_tests.MLXTestCase):
             with self.subTest(dtype=dtype, axis=axis, strided=strided):
                 np.random.seed(0)
                 np_dtype = getattr(np, dtype)
-                a_np = np.random.uniform(0, 100, size=shape).astype(np_dtype)
+                if np.issubdtype(np_dtype, np.complexfloating):
+                    a_np = (
+                        np.random.uniform(0, 100, size=shape)
+                        + 1j * np.random.uniform(0, 100, size=shape)
+                    ).astype(np_dtype)
+                else:
+                    a_np = np.random.uniform(0, 100, size=shape).astype(np_dtype)
                 a_mx = mx.array(a_np)
                 if strided:
                     a_mx = a_mx[::2, :, ::2]
@@ -3317,7 +3326,10 @@ 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))
 
-        x = mx.array([3.0, mx.nan, 2.0, 0.0]) + 1j * mx.array([1.0] * 4)
+        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))
 
     def test_argsort_nan(self):
         for dtype in [mx.float32, mx.float16, mx.bfloat16]: