Skip to content

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

Description

@mateuuszzzzz

☑️ 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: #4529

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions