[BUG] Metal `col_reduce_longcolumn` produces wrong results for negative-stride views
☑️ 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
Source: ml-explore/mlx