[BUG] `argsort` diverges between cpu and gpu on complexes with NaN
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 28.5k
- Forks
- 2.3k
- Avg merge
- 3d 8h
- Merged PRs (30d)
- 62
Description
Describe the bug
An array with complex64 values containing NaNs sorts differently on gpu and cpu streams.
To Reproduce
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]
GPU streams behave correctly (input is sorted, NaN last, result matches np.argsort)
CPU streams return [0,1,2,...n] without sorting.
Expected behavior
Expected outputs to match. Ideally they should behave as numpy et al.
Desktop (please complete the following information):
- OS Version: MacOS 27.0 RC
- Version :
main(v0.32.2 + d9add9d11)
Additional context
Hit in PR #4501
Issue 100% written by me, a human.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reproducing the complex64 NaN case with the Python mx.argsort entry point on both mx.cpu and mx.gpu streams. Trace the CPU and GPU argsort paths, then add or update coverage for the shown input so both outputs match numpy with NaN last; the issue is done when CPU and GPU results agree.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- data
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 68/100