ml-explore / ml-explore/mlx

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

Open
#4,502 0 comments 0 reactions 0 assignees View on GitHub

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.