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 + d9add9d)
Additional context
Hit in PR #4501
Issue 100% written by me, a human.
Describe the bug
An array with
complex64values containing NaNs sorts differently on gpu and cpu streams.To Reproduce
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):
main(v0.32.2 + d9add9d)Additional context
Hit in PR #4501
Issue 100% written by me, a human.