From ef23f4b6ee3ccecdf9647527bc2135a731a34118 Mon Sep 17 00:00:00 2001 From: Roberto Meroni Date: Thu, 13 Aug 2026 09:22:38 +0200 Subject: [PATCH] Fix Metal sort of a view with a negative stride The non-contiguous sort kernels compute each row's base offset with elem_to_loc(tid.y, ...). IdxT is deduced from the argument, so tid.y makes it uint and a negative stride wraps to ~4.29e9 instead of stepping backwards, so every row but the first reads out of bounds and Metal returns zeros. Ask for int64_t, matching what the CUDA kernels already do with int64_t(blockIdx.y). --- mlx/backend/metal/kernels/sort.h | 10 +++++++--- python/tests/test_ops.py | 21 +++++++++++++++++++++ 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/mlx/backend/metal/kernels/sort.h b/mlx/backend/metal/kernels/sort.h index 068d43d126..ea2640bace 100644 --- a/mlx/backend/metal/kernels/sort.h +++ b/mlx/backend/metal/kernels/sort.h @@ -388,8 +388,11 @@ template < using ValT = typename sort_kernel::ValT; using IdxT = typename sort_kernel::IdxT; - auto in_block_idx = elem_to_loc(tid.y, nc_shape, in_nc_strides, nc_dim); - auto out_block_idx = elem_to_loc(tid.y, nc_shape, out_nc_strides, nc_dim); + // Signed offsets: a non-sorted axis may have a negative stride. + auto in_block_idx = + elem_to_loc(tid.y, nc_shape, in_nc_strides, nc_dim); + auto out_block_idx = + elem_to_loc(tid.y, nc_shape, out_nc_strides, nc_dim); inp += in_block_idx; out += out_block_idx; @@ -532,7 +535,8 @@ template < BLOCK_THREADS, N_PER_THREAD>; - auto block_idx = elem_to_loc(tid.y, nc_shape, nc_strides, nc_dim); + // Signed offset: a non-sorted axis may have a negative stride. + auto block_idx = elem_to_loc(tid.y, nc_shape, nc_strides, nc_dim); inp += block_idx; out_vals += tid.y * size_sorted_axis; out_idxs += tid.y * size_sorted_axis; diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index d569dbd3c7..d5ae813efe 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2675,6 +2675,27 @@ def test_sort(self): y_np = np.sort(np.array(a), axis=-1) self.assertTrue(np.array_equal(y_np, y_mx)) + # Negative stride on an axis that is not sorted, single and multi block + np.random.seed(0) + for dtype in ("int32", "float32"): + for size in (4, 32769): + with self.subTest(dtype=dtype, size=size): + a_np = np.random.uniform(0, 100, size=(3, size)) + a_np = a_np.astype(getattr(np, dtype)) + a_mx = mx.array(a_np)[::-1, :] + a_np = a_np[::-1, :] + + b_np = np.sort(a_np, axis=-1) + self.assertTrue(np.array_equal(b_np, mx.sort(a_mx, axis=-1))) + + idx = mx.argsort(a_mx, axis=-1) + self.assertTrue( + np.array_equal(b_np, mx.take_along_axis(a_mx, idx, axis=-1)) + ) + + b_mx = mx.partition(a_mx, 1, axis=-1) + self.assertTrue(np.array_equal(b_np[:, 1], np.array(b_mx)[:, 1])) + def test_partition(self): shape = (3, 4, 5) for dtype in ("int32", "float32"):