Skip to content

[Metal] Support sorting complex numbers#3314

Open
Lyxot wants to merge 2 commits intoml-explore:mainfrom
Lyxot:metal-complex-sort
Open

[Metal] Support sorting complex numbers#3314
Lyxot wants to merge 2 commits intoml-explore:mainfrom
Lyxot:metal-complex-sort

Conversation

@Lyxot
Copy link
Copy Markdown
Contributor

@Lyxot Lyxot commented Mar 25, 2026

Proposed changes

Followup to #3286.

Add complex64 sort, argsort, partition, and argpartition support to the Metal backend

Copilot AI review requested due to automatic review settings March 25, 2026 06:51
Copy link
Copy Markdown

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds Metal backend support for sorting complex64 values (including NaN-aware ordering) and updates Python tests to exercise complex sorting on Metal, as a follow-up to CUDA complex sorting support.

Changes:

  • Instantiate Metal block and multi-block sort/partition/merge kernels for complex64.
  • Add complex64_t NaN initialization and isnan(complex64_t) support for comparisons in sort kernels.
  • Expand Python sort tests to include complex64 unconditionally (including NaN behavior).

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File Description
python/tests/test_ops.py Enables complex64 coverage for sort and complex NaN sort expectations on Metal
mlx/backend/metal/kernels/sort.metal Instantiates sort-related kernels for complex64 in the static Metal kernel set
mlx/backend/metal/kernels/sort.h Adds complex64 init value (NaN) and enables NaN-aware comparator behavior for complex64
mlx/backend/metal/kernels/complex.h Adds isnan(complex64_t) helper used by sort comparator

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copy link
Copy Markdown
Collaborator

@zcbenz zcbenz left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me!

@Lyxot Lyxot force-pushed the metal-complex-sort branch 3 times, most recently from 075ee43 to 8944273 Compare March 25, 2026 13:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants