[WIP] Add mlx.core.searchsorted #2817
Draft
Conversation
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
Introduces a comprehensive Product Requirements Document for the new mlx.searchsorted feature, outlining API parity with NumPy, multi-dimensional support, axis-based operations, and GPU-accelerated implementation plans. Documents goals, functional/spec requirements, testing strategies, documentation plans, release notes, risks, and implementation phases. References the related issue ml-explore#1255.
Introduces a comprehensive test suite for the searchsorted operation, validating correctness across multiple dimensions, dtypes (int, float, half precision), and sides, including: - Basic 1D behavior and dtype compatibility - Handling of duplicates and edge cases (empty, single element, all duplicates) - Multidimensional inputs with various axes and broadcasting rules - Special values (inf, -inf, NaN) and out-of-range scenarios - Output shape and dtype correctness for scalar, vector, and 2D value inputs - Mixed precision scenarios to ensure robust type handling - Large input arrays for performance/behavior checks
Introduces a .clangd configuration to enhance editor tooling and code quality. - Enables strict diagnostic checks (UnusedIncludes, MissingIncludes) and enables Clang-Tidy with tailored rule sets. - Configures compilation database and background indexing to speed up tooling feedback. - Enables inlay hints (parameter names, deduced types) and hover improvements for better readability. - Activates comprehensive completion across all scopes and improves developer experience with AKA hover support.
Drops unused header dependencies: - removes transforms.h from the operations module - removes backend/utils.h from primitives
Introduces a local numpy-based reference implementation for searchsorted that supports axis and broadcasting, intended to validate the multi-dimensional behavior against the library. Updates test expectations for scalar input on 2D arrays. Temporarily comments out extensive multi-dimensional test cases pending full parity, but leaves the groundwork for future validation and comparison.
Improves code hygiene by: - Wrapping a long `set` command in `CMakeLists.txt` for better readability. - Removing trailing whitespace in `test_ops.py`. - Standardizing string literal quotes (single to double) in `test_ops.py`. - Converting whitespace-only lines to truly empty lines.
This reverts commit abf4f49.
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.
First time contributor, appreciate any feedback! Work in progress on feature req from #1255.
Proposed changes
In Progress:
[ ] Add functionality for axis parameter for parity with np.searchsorted (documentation has it currently)
[ ] Run benchmarks
[ ] (not optimized) Implement vectorized linear search vs non-vectorized binary search
[ ] implement eval_gpu.
Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes