Skip to content

Fix the CPU scan over a size one axis with a padded stride - #4139

Draft
kapellirohith wants to merge 1 commit into
ml-explore:mainfrom
kapellirohith:cpu-scan-size-one-axis
Draft

Fix the CPU scan over a size one axis with a padded stride#4139
kapellirohith wants to merge 1 commit into
ml-explore:mainfrom
kapellirohith:cpu-scan-size-one-axis

Conversation

@kapellirohith

Copy link
Copy Markdown
Contributor

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:

x = mx.arange(1, 11, dtype=mx.float32).reshape(1, 10)[:, 2:]
mx.cumsum(x, axis=0)

The strided CPU scan takes its row count from
size / shape[axis] / strides[axis]. row_contiguous ignores 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 count
floors 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 --check against
pristine main at scan.cpp:200, independently of this change. The resolution
keeps both its AccT init and the || in.shape(axis) == 1 here. The two are
semantically 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.

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
kapellirohith marked this pull request as draft August 10, 2026 16:39
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.

1 participant