From ba7b28fac91666cd3d19f1c7e3628c978e9f6205 Mon Sep 17 00:00:00 2001 From: Aditya Jain Date: Thu, 27 Aug 2026 01:41:24 -0700 Subject: [PATCH] Fix setitem negative index normalized against wrong axis when preceded by None mlx_compute_slice_update_args normalized a negative scalar index using src.shape(i), the position in the raw index tuple, instead of src.shape(ax), the actual source axis. These diverge whenever a None (newaxis) precedes the scalar index, since None consumes a tuple slot but not a source axis. The wrong axis size produced an out-of-range start/stop for the slice_update, which silently no-ops instead of writing anything. Example: a[None, -1] = 9 on a (3, 4) array wrote nothing at all, instead of updating the last row. --- python/src/indexing.cpp | 2 +- python/tests/test_array.py | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/python/src/indexing.cpp b/python/src/indexing.cpp index ce64b7bd3e..0e230ba529 100644 --- a/python/src/indexing.cpp +++ b/python/src/indexing.cpp @@ -913,7 +913,7 @@ mlx_compute_slice_update_args( upd_ax--; } else if (is_index_scalar(pyidx)) { int st = safe_to_int32(pyidx); - st = (st < 0) ? st + src.shape(i) : st; + st = (st < 0) ? st + src.shape(ax) : st; starts[ax] = st; stops[ax] = st + 1; if (upd_ax >= 0) { diff --git a/python/tests/test_array.py b/python/tests/test_array.py index ff9343e290..f82ece534d 100644 --- a/python/tests/test_array.py +++ b/python/tests/test_array.py @@ -1348,6 +1348,13 @@ def test_setitem(self): a[0:1] = mx.array([1]) self.assertEqual(a.tolist(), [1, 4, 4]) + # Regression test: a negative integer index after a None + # (newaxis) used to be normalized against the wrong axis size, + # silently writing nothing instead of updating the last row. + b = mx.zeros((3, 4)) + b[None, -1] = 9 + self.assertEqual(b.tolist(), [[0, 0, 0, 0], [0, 0, 0, 0], [9, 9, 9, 9]]) + with self.assertRaises(ValueError): a[0:1] = mx.array([2, 3])