#4502·mlx

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

Author: louenCreated Sep 14, 2026Updated Sep 14, 2026

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.