Skip to content

Fix ordering of complex64 with NaNs - #4519

Open
louen wants to merge 1 commit into
ml-explore:mainfrom
louen:val/fix-cpu-complex-sort-nan
Open

louen wants to merge 1 commit into
ml-explore:mainfrom
louen:val/fix-cpu-complex-sort-nan

Conversation

@louen

@louen louen commented Sep 15, 2026

Copy link
Copy Markdown
Contributor

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

@louen
louen marked this pull request as ready for review September 15, 2026 23:04
@louen louen changed the title std::isnan to mlx::isnan Fix ordering of complex64 with NaNs Sep 15, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] argsort diverges between cpu and gpu on complexes with NaN

2 participants