[BUG] Metal `col_reduce_longcolumn` produces wrong results for negative-stride views
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 28.5k
- Forks
- 2.3k
- Avg merge
- 3d 8h
- Merged PRs (30d)
- 62
Description
☑️ I understand it is strictly prohibited to use AI to write issues.
Describe the bug
On the Metal backend, reductions return wrong results for some views with a negative stride. The CPU backend and NumPy give the correct result (same goes for the same array without the reversal).
The snippet uses sum, but the same happens for prod, max, min, all, any, and also mean and var, which are derived from sum.
Shape impacts correctness, only one of the column-reduction kernels is affected, and other shapes e.g.(2, 512, 64) control in the snippet, are dispatched to kernels that work correctly.
To Reproduce
import mlx.core as mx
import numpy as np
print(mx.__version__)
x = mx.arange(1, 2 * 1024 * 16 + 1).reshape(2, 1024, 16)[::-1]
x_np = np.arange(1, 2 * 1024 * 16 + 1).reshape(2, 1024, 16)[::-1]
print("GPU: ", mx.sum(x, axis=1, stream=mx.gpu)[1, :4])
print("CPU: ", mx.sum(x, axis=1, stream=mx.cpu)[1, :4])
print("NumPy:", x_np.sum(axis=1)[1, :4])
# same shape without the reversed view: correct
x = mx.arange(1, 2 * 1024 * 16 + 1).reshape(2, 1024, 16)
print(
"plain array, GPU == CPU:",
mx.array_equal(mx.sum(x, axis=1, stream=mx.gpu), mx.sum(x, axis=1, stream=mx.cpu)).item(),
)
# reversed view, but last dim 64 (a different column kernel): correct
y = mx.arange(1, 2 * 512 * 64 + 1).reshape(2, 512, 64)[::-1]
print(
"last dim 64, GPU == CPU:",
mx.array_equal(mx.sum(y, axis=1, stream=mx.gpu), mx.sum(y, axis=1, stream=mx.cpu)).item(),
)
Output:
0.32.2
GPU: array([0, 0, 0, 0], dtype=int32)
CPU: array([8381440, 8382464, 8383488, 8384512], dtype=int32)
NumPy: [8381440 8382464 8383488 8384512]
plain array, GPU == CPU: True
last dim 64, GPU == CPU: True
Expected behavior
Metal backend should return the same results as NumPy and the MLX CPU backend.
Desktop (please complete the following information):
- OS Version: macOS 26.6.2
- Version: 0.32.2
Additional context
The bug needs all three of the following at once:
- a negative stride on a non-reduced axis before the reduced one e.g.
x[::-1]reduced overaxis=1forxof shape(2, 1024, 16)(a negative stride after the reduced axis picks the contiguous-copy flow) - a reduction over a non-last axis
- dispatch to
col_reduce_longcolumn(reduction_stride < 32andreduction_size * non_col_reductions >= 1024)
Proposed fix: https://github.com/ml-explore/mlx/pull/4529
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the Metal col_reduce_longcolumn kernel and reproduce the issue using the Python snippet in the report. Compare the reduction results with the CPU backend and NumPy for the reversed view, then verify that sum, prod, max, min, all, and any produce matching results across the affected shape; the issue links proposed fix PR #4529.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- backend, performance
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 25/100