Fix the CPU scan over a size one axis with a padded stride - #4139
Draft
kapellirohith wants to merge 1 commit into
Draft
Fix the CPU scan over a size one axis with a padded stride#4139kapellirohith wants to merge 1 commit into
kapellirohith wants to merge 1 commit into
Conversation
The strided scan takes its row count from size / shape[axis] / strides[axis]. A size one axis is allowed to carry any stride and still be row contiguous, so for a slice such as x[:, 2:] of a (1, N) array the stride is larger than the number of elements and the count floors to zero. The scan then wrote nothing and the output kept whatever was in the freshly allocated buffer. A padded stride is only possible on a size one axis, and the scan over such an axis is elementwise, which is what the contiguous path already computes.
kapellirohith
marked this pull request as draft
August 10, 2026 16:39
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
A scan over a size one axis that carries a padded stride leaves the output
unwritten, so it returns whatever was in the freshly allocated buffer:
The strided CPU scan takes its row count from
size / shape[axis] / strides[axis].row_contiguousignores size one axes,so such an axis can carry any stride and the array is still flagged row
contiguous. Here the shape is
(1, 8)with strides(10, 1), the countfloors to zero, and the loop body never runs.
A padded stride is only possible on a size one axis, and a scan over a size
one axis is elementwise, which is what the contiguous path with a stride of
one already computes, so no copy is needed.
Introduced by c423074 "redesign for faster cpu/gpu synch (#1869)"
(2025-03-06).
The GPU backends are not affected. Both Metal and CUDA give the output the
input's strides and flags, so the two stay in the same frame, and their row
count comes from
data_size()rather than the stride.Note for whoever lands second: #3907 already fails
git apply --checkagainstpristine main at
scan.cpp:200, independently of this change. The resolutionkeeps both its
AccT initand the|| in.shape(axis) == 1here. The two aresemantically orthogonal.
Ran locally on an M3 Pro: the Python suite on the GPU and with
mx.set_default_device(mx.cpu), the C++ suite on both devices, pre-commit,and the new test five times on each device. Not covered locally: CUDA, Linux,
Windows.