Fix ordering of complex64 with NaNs

[louen]

This PR fixes #4502

Bug

The symptom was that complex types with NaNs in their imaginary part are sorted differently on CPU and GPU:

import mlx.core as mx
NAN = float("nan")
a = mx.array([complex(1, NAN), 1 + 1j, 2 + 0j], mx.complex64)
mx.argsort(a, stream=mx.cpu) # -> [0 1 2]
mx.argsort(a, stream=mx.gpu) # -> [1 2 0]

The reason is in the CPU implementation of sort.cpp, which uses std::isnan instead of mlx::core::isnan.
This is problematic because our mlx::complex type has a cast operator() float operator which returns the real part.
Therefore, using std::isnan only returns true if the real part is a NaN, but ignores the imaginary part, hence the bug.
The GPU implementation did not suffer from this issue, e.g. Metal sort has explicit isnan for both parts

 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);   

Fix

The fix was first to convert current std::isnan to mlx::core::isnan for all arguments that could be either floating point types or complex.

A better fix would be to make the operator explicit, which would prevent silent cast the complex type to float but that seems a larger impact.

Verified that the fix works as expected:

%> pip install -e ".[dev]"  && python3
>>> import mlx.core as mx
>>> NAN = float("nan")
>>> a = mx.array([complex(1, NAN), 1 + 1j, 2 + 0j], mx.complex64)
>>> mx.argsort(a, stream=mx.cpu) 
array([1, 2, 0], dtype=uint32)
>>> mx.argsort(a, stream=mx.gpu)
array([1, 2, 0], dtype=uint32)
>>> 

• AI usage disclosure: Claude Code (Opus 5) checked for coverage of std::isnan usage on possible complex args, and produced the tests.
• ☑️ I understand it is strictly prohibited to use AI to write PR description

[RohanGautam]

Good catch! Makes sense to me.

[RohanGautam]

Can you rebase on main? Some tests are failing which were later fixed

[louen]

Branch rebased