diff --git a/mlx/backend/metal/kernels/reduction/reduce_row.h b/mlx/backend/metal/kernels/reduction/reduce_row.h index 936d75bb52..b55c83f315 100644 --- a/mlx/backend/metal/kernels/reduction/reduce_row.h +++ b/mlx/backend/metal/kernels/reduction/reduce_row.h @@ -337,7 +337,8 @@ template < // lid.x * N_READS breaks the per_thread_row_reduce interface a bit. Maybe it // needs a small refactor. - in += elem_to_loc(out_idx, shape, strides, ndim) + lid.x * N_READS; + in += + elem_to_loc(out_idx, shape, strides, ndim) + IdxT(lid.x) * N_READS; LoopedElemToLoc 2)> loop(reduce_ndim); const device T* row; diff --git a/python/tests/test_reduce.py b/python/tests/test_reduce.py index 6ac8fc1504..164e2dd803 100644 --- a/python/tests/test_reduce.py +++ b/python/tests/test_reduce.py @@ -1,6 +1,5 @@ # Copyright © 2023 Apple Inc. -import unittest from itertools import combinations, permutations import mlx.core as mx @@ -47,6 +46,16 @@ def test_expand_sums(self): np.allclose(z_npy, np.array(z_mlx), atol=1e-4) ) + def test_row_reduce_negative_stride(self): + x_npy = np.arange(1, 131).reshape(2, 65)[::-1] + x_mlx = mx.arange(1, 131).reshape(2, 65)[::-1] + + for op in ["sum", "max", "min", "mean", "var"]: + with self.subTest(op=op): + expected = getattr(np, op)(x_npy, axis=-1) + actual = getattr(mx, op)(x_mlx, axis=-1) + self.assertTrue(np.allclose(expected, actual)) + def test_dtypes(self): int_dtypes = [ "int8",