Repository navigation
Fix segfault in expand_dims for out of bounds negative axes - #4021
Merged
Merged
Conversation
expand_dims(a, axes) normalized negative axes against the output rank and then passed them to expand_dims_impl, which normalized them a second time. A negative axis that was still negative after the first pass was silently mapped to a valid but wrong axis, and the resulting axes were no longer sorted. ExpandDims::output_shape assumes sorted in bounds axes, so mx.expand_dims(mx.zeros(()), (-3, -2)) inserted past the end of the shape and crashed. Normalize, validate and sort the axes once via normalize_axis_index. The same double normalization was present in squeeze, where it silently squeezed the wrong axis instead of crashing.
Adityaj0
added a commit
to Adityaj0/mlx
that referenced
this pull request
Aug 10, 2026
…ong_axis
The bounds check in take_along_axis, scatter_axis (which backs put_along_axis
and scatter_add_axis) and linalg::cross was dead code:
if (axis + a.ndim() < 0 || axis >= static_cast<int>(a.ndim()))
array::ndim() returns size_t, so `axis + a.ndim()` is evaluated in unsigned
arithmetic and is never negative. Only the positive half of the check, which
casts to int, actually ran. A negative out of bounds axis therefore flowed
through to broadcast_arrays as `axis - int(a.ndim())` and either raised an
internal "SmallVector out of range" error or, for larger magnitudes,
overwrote the stack and segfaulted inside broadcast_shapes.
Use normalize_axis_index, which does the comparison in int and is already
used for this purpose elsewhere in ops.cpp. It also normalizes, so the later
manual "allow negative axis" adjustments are no longer needed.
Same class of bug as ml-explore#4021.
Co-Authored-By: Claude Opus 5 <[email protected]>
3 of 4 tasks
Adityaj0
added a commit
to Adityaj0/mlx
that referenced
this pull request
Aug 10, 2026
…ong_axis
The bounds check in take_along_axis, scatter_axis (which backs put_along_axis
and scatter_add_axis) and linalg::cross was dead code:
if (axis + a.ndim() < 0 || axis >= static_cast<int>(a.ndim()))
array::ndim() returns size_t, so `axis + a.ndim()` is evaluated in unsigned
arithmetic and is never negative. Only the positive half of the check, which
casts to int, actually ran. A negative out of bounds axis therefore flowed
through to broadcast_arrays as `axis - int(a.ndim())` and either raised an
internal "SmallVector out of range" error or, for larger magnitudes,
overwrote the stack and segfaulted inside broadcast_shapes.
Use normalize_axis_index, which does the comparison in int and is already
used for this purpose elsewhere in ops.cpp. It also normalizes, so the later
manual "allow negative axis" adjustments are no longer needed.
Same class of bug as ml-explore#4021.
1 task done
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
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.
Proposed changes
Closes #4011
mx.expand_dims(mx.zeros(()), axis=(-3, -2))segfaults.Cause
expand_dims(a, axes)normalized negative axes against the output rank andthen passed them to
expand_dims_impl, which normalized them a second time(
-3 + 2 = -1, still negative and never range checked, then-1 + 2 = 1). Anout of bounds axis therefore turned silently into a valid but wrong one, and
the axes ended up as
{1, 0}instead of sorted.ExpandDims::output_shapeassumes sorted, in bounds axes, so this inserted past the end of an empty shape
and crashed.
out_ndimwas alsosize_t, so thenew_ax < 0guard couldnever fire.
Fix
Normalize, validate and sort the axes exactly once, in a single helper per op,
reusing the existing
normalize_axis_index. Error wording now comes from thathelper; no tests asserted on the old strings.
squeezehad the same double normalization, but failed silently rather thancrashing:
mx.squeeze(mx.zeros((1, 1, 1)), (-5,))returned shape(1, 1)instead of raising, because
-5mapped to axis1. Fixed the same way, thoughhappy to split it out if you would prefer this PR stay scoped to the crash.
Testing
Extended the C++
test squeeze and expandcase and addedtest_squeeze_expand_invalid_axesin Python, covering out of bounds negativeaxes for both ops plus the valid negative axis cases that must keep working.
Both segfault against unpatched
ops.cppand pass with the fix.Full C++ and Python suites pass on a Metal build; the only failures are torch
MPS DLPack tests, which fail identically without this change.
Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes