Skip to content

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

Description

@louen

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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions