Skip to content

Fix Metal row reductions on negative-stride views - #4267

Merged
zcbenz merged 2 commits into
ml-explore:mainfrom
FU-max-boop:fix/metal-negative-stride-row-reduction
Aug 16, 2026
Merged

Fix Metal row reductions on negative-stride views#4267
zcbenz merged 2 commits into
ml-explore:mainfrom
FU-max-boop:fix/metal-negative-stride-row-reduction

Conversation

@FU-max-boop

@FU-max-boop FU-max-boop commented Aug 15, 2026

Copy link
Copy Markdown
Contributor

Proposed changes

Metal row reductions return incorrect values when a non-reduced dimension has a negative stride and the reduced row contains more than 64 elements. For example:

x = mx.arange(1, 131).reshape(2, 65)[::-1]
mx.sum(x, axis=-1, stream=mx.gpu)

NumPy and the MLX CPU backend return [6370, 2145]; current Metal returns [6370, 0]. The same pointer error affects max, min, mean, and var.

row_reduce_looped combines the signed offset returned by elem_to_loc<IdxT> with the unsigned expression lid.x * N_READS. For the ordinary IdxT=int specialization, a negative outer-dimension offset is therefore promoted to uint and wraps before pointer arithmetic. Casting lid.x to IdxT keeps this addition signed. The large precompiled IdxT=int64_t specialization is unchanged semantically.

This adds an explicit Metal regression test for five row reductions on a 65-element negative-stride view. This is adjacent to the signed-indexing issue addressed for sort kernels in #4252, but affects the separate row-reduction kernel.

Validation

  • Current main (140faa8a) / development Metal wheel: the new regression test fails for all five operations.
  • Patched source, CPU-only build: python -m unittest test_reduce — 13 passed; the Metal-only regression is skipped as intended.
  • pre-commit run --all-files — passed.
  • git diff --check — passed.

This machine has the Command Line Tools but not the standalone Xcode Metal compiler, so the patched metallib has not been rebuilt locally. The repository's macOS CI must compile and validate the Metal change before merge.

Disclosure: Codex assisted with differential testing and review. I verified the reproduction, root cause, focused diff, and local checks before submission.

Checklist

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (not needed; no API change)

@FU-max-boop
FU-max-boop marked this pull request as ready for review August 15, 2026 09:46
Comment thread python/tests/test_reduce.py Outdated

@zcbenz zcbenz left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Thanks!

@zcbenz
zcbenz merged commit a1e0e0b into ml-explore:main Aug 16, 2026
28 checks passed
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.

2 participants