From 5b8c7dc92d3a79a0b5fc4035f3bfcc8b82bcc2bd Mon Sep 17 00:00:00 2001 From: Mateusz Rajski Date: Thu, 17 Sep 2026 21:28:00 +0200 Subject: [PATCH 1/2] Fix Metal col_reduce_longcolumn for negative-stride views --- mlx/backend/metal/kernels/reduction/reduce_col.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx/backend/metal/kernels/reduction/reduce_col.h b/mlx/backend/metal/kernels/reduction/reduce_col.h index b1546adb55..5dfc624bb4 100644 --- a/mlx/backend/metal/kernels/reduction/reduce_col.h +++ b/mlx/backend/metal/kernels/reduction/reduce_col.h @@ -117,7 +117,7 @@ template IdxT out_idx = gid.x + gsize.x * IdxT(gid.y); IdxT in_idx = elem_to_loc(out_idx, shape, strides, ndim); - in += in_idx + lid.x; + in += in_idx + IdxT(lid.x); U total = Op::init; IdxT total_rows = IdxT(non_col_reductions) * IdxT(reduction_size); From 54043df76605e67fcfad783509a69a60141ca65b Mon Sep 17 00:00:00 2001 From: Mateusz Rajski Date: Thu, 17 Sep 2026 22:11:28 +0200 Subject: [PATCH 2/2] Add tests covering col_reduce_* for negative-stride views --- python/tests/test_reduce.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/python/tests/test_reduce.py b/python/tests/test_reduce.py index 164e2dd803..0b04038176 100644 --- a/python/tests/test_reduce.py +++ b/python/tests/test_reduce.py @@ -1,5 +1,6 @@ # Copyright © 2023 Apple Inc. +import math from itertools import combinations, permutations import mlx.core as mx @@ -56,6 +57,33 @@ def test_row_reduce_negative_stride(self): actual = getattr(mx, op)(x_mlx, axis=-1) self.assertTrue(np.allclose(expected, actual)) + def test_col_reduce_negative_stride(self): + # Exercises each Metal column reduction kernel on a negative-stride view. + cases = [ + ((2, 1024, 16), 1), + ((2, 2, 1024, 16), 2), + ((2, 512, 2, 2, 16), (1, 3)), + ((2, 512, 64), 1), + ((2, 64, 512), 1), + ((2, 16, 16), 1), + ] + for shape, axis in cases: + size = math.prod(shape) + x_npy = np.arange(1, size + 1).reshape(shape)[::-1] + x_mlx = mx.arange(1, size + 1).reshape(shape)[::-1] + for op in ["sum", "max", "min", "mean", "var"]: + with self.subTest(shape=shape, axis=axis, op=op): + expected = getattr(np, op)(x_npy, axis=axis) + actual = getattr(mx, op)(x_mlx, axis=axis) + self.assertTrue(np.allclose(expected, actual)) + x_ones = mx.ones(shape, dtype=mx.float32)[::-1] + with self.subTest(shape=shape, axis=axis, op="prod"): + self.assertTrue(mx.all(mx.prod(x_ones, axis=axis) == 1).item()) + x_bool = mx.ones(shape, dtype=mx.bool_)[::-1] + for op in ["all", "any"]: + with self.subTest(shape=shape, axis=axis, op=op): + self.assertTrue(mx.all(getattr(mx, op)(x_bool, axis=axis)).item()) + def test_dtypes(self): int_dtypes = [ "int8",