Document sort stability and NaN handling#3400
Open
NeuralNoble wants to merge 1 commit intoml-explore:mainfrom
Open
Document sort stability and NaN handling#3400NeuralNoble wants to merge 1 commit intoml-explore:mainfrom
NeuralNoble wants to merge 1 commit intoml-explore:mainfrom
Conversation
Both `sort` and `argsort` are stable across all backends (std::stable_sort on CPU, merge sort on Metal/CUDA) and place NaN values at the end to match NumPy/JAX. Add this to the Python docstrings and C++ declarations so users can rely on these guarantees.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Addresses #2057.
Documents that
sortandargsortare stable across all backends and place NaN values at the end to match NumPy/JAX.std::stable_sortwith a NaN-aware comparator (mlx/backend/cpu/sort.cpp)mlx/backend/metal/kernels/sort.h,mlx/backend/cuda/sort.cu)Both guarantees are true in the current implementation (NaN handling introduced in #2667) but weren't documented, so users couldn't rely on them.
Changes
python/src/ops.cpp— add note tosortandargsortdocstringsmlx/ops.h— add note to the four C++ declarationsmlx/ops.cpp— add note to the four C++ definitionsTest plan
pre-commit run --all-filespassesctestpassespython -m pytest python/tests/test_ops.py -k sortpasseshelp(mx.sort)renders the new docstring at runtime