ml-explore / ml-explore/mlx

[BUG] Metal `col_reduce_longcolumn` produces wrong results for negative-stride views

Open
#4,528 0 comments 0 reactions 0 assignees View on GitHub

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:

  1. a negative stride on a non-reduced axis before the reduced one e.g. x[::-1] reduced over axis=1 for x of shape (2, 1024, 16) (a negative stride after the reduced axis picks the contiguous-copy flow)
  2. a reduction over a non-last axis
  3. dispatch to col_reduce_longcolumn (reduction_stride < 32 and reduction_size * non_col_reductions >= 1024)

Proposed fix: https://github.com/ml-explore/mlx/pull/4529

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.