Pass axis as a Python int in the MLX sort/argsort dispatches - #2400
Open
guillaume-osmo wants to merge 2 commits into
Open
Pass axis as a Python int in the MLX sort/argsort dispatches#2400guillaume-osmo wants to merge 2 commits into
guillaume-osmo wants to merge 2 commits into
Conversation
`mlx_funcify_Reshape` forwarded the shape input straight to `mx.reshape`, but the linker typifies every input to `mx.array` while `mx.reshape` only accepts a Python sequence of ints, so every reshape raised `TypeError`. Reading the array back at runtime is not an option either: the linker enables `mx.compile` by default and MLX forbids evaluating a traced array, so the shape has to be resolved when the dispatch is built. Prefer the statically inferred output shape (which already resolves any `-1`, and covers shapes read off another input such as `x.reshape(y.shape)`), then fall back to a constant shape input. A genuinely data-dependent shape now raises `NotImplementedError` with an explanation instead of an eval error from deep inside MLX. This turns four already-red tests in `tests/link/mlx/test_shape.py` green and adds regression tests under the full "MLX" mode.
`mlx_funcify_Sort` and `mlx_funcify_ArgSort` forwarded the `axis` input straight to `mx.sort` / `mx.argsort`, but the linker typifies every input to `mx.array` and both functions require a Python int, so nothing on this backend sorted at all -- every call raised `TypeError`. Same root cause as the `Reshape` dispatch (pymc-devs#2386), which `axis=None` also goes through, since it flattens first. All five tests in `tests/link/mlx/test_sort.py` were already red on main; they pass now, and a parametrized test covers positive, negative and `None` axes on a 3-d input. Found while fixing pymc-devs#2385: the gradient of a batched `solve` reaches `argsort` through the pivot-to-permutation step.
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.
Closes #2391.
Motivation
mlx_funcify_Sortandmlx_funcify_ArgSortforwarded theaxisinput straight tomx.sort/mx.argsort, but the linker typifies every input tomx.arrayand both require a Pythonint, so nothing on this backend sorted at all. Same root cause as #2386.All five tests in
tests/link/mlx/test_sort.pywere already red onmain.Found while fixing #2385: the gradient of a batched
solvereachesargsortthrough the pivot-to-permutation step.Implementation
Coerce
axiswithint(...). Unlike the shape in #2386,axisis a scalar and always static, so no funcify-time resolution is needed.Tests
The five existing tests pass now. Added
test_sort_axis_variants, covering positive, negative andNoneaxes for bothsortandargsorton a 3-d input.Full
tests/link/mlx/suite: no regressions.