From 7c51cc55f558705950faea3b71c4eb7ace2e5dc3 Mon Sep 17 00:00:00 2001 From: guillaume-osmo Date: Wed, 26 Aug 2026 15:53:07 +0200 Subject: [PATCH] Honour idx_list in the MLX AdvancedIncSubtensor dispatch (#2382, #2387) `mlx_funcify_AdvancedIncSubtensor` passed the flat index list straight to the scatter, ignoring `op.idx_list`. `idx_list` records where the index variables sit among plain slices, so dropping it anchored the advanced indices at axis 0 and shifted every one of them left. Any destination with leading batch axes then failed to broadcast: idx_list=(slice(None), 0) -> #2382, [broadcast_shapes] (4,1) vs (7,3) idx_list=(slice(None), 0, 1) -> #2387, [broadcast_shapes] (5,3) vs (3,3) Both are easy to hit without writing any indexing, since `specialize` rewrites `diagonal(...)` into this form -- the gradient of a log-determinant over a batch of matrices was enough. Call `unflatten_index_variables`, mirroring `AdvancedIncSubtensor.perform`, which is also what the `AdvancedSubtensor` dispatch already does. The `set_instead_of_inc` and `ignore_duplicates` branches take the same path, and MLX accepts a slice in the index tuple for all three. --- pytensor/link/mlx/dispatch/subtensor.py | 14 +++++- tests/link/mlx/test_subtensor.py | 66 +++++++++++++++++++++++++ 2 files changed, 78 insertions(+), 2 deletions(-) diff --git a/pytensor/link/mlx/dispatch/subtensor.py b/pytensor/link/mlx/dispatch/subtensor.py index 6fded559df..8a58f65f01 100644 --- a/pytensor/link/mlx/dispatch/subtensor.py +++ b/pytensor/link/mlx/dispatch/subtensor.py @@ -7,6 +7,7 @@ IncSubtensor, Subtensor, indices_from_subtensor, + unflatten_index_variables, ) @@ -96,9 +97,18 @@ def mlx_fn(x, indices, y): def mlx_fn(x, indices, y): return x.at[indices].add(y) - def advancedincsubtensor(x, y, *ilist, mlx_fn=mlx_fn): + idx_list = op.idx_list + + def advancedincsubtensor(x, y, *ilist, mlx_fn=mlx_fn, idx_list=idx_list): op._check_runtime_broadcast_of_vector_index(node, x, y, ilist[0]) - return mlx_fn(x, ilist, y) + # `idx_list` records where the index variables sit among plain slices, + # so the advanced indices are not necessarily anchored at axis 0. Using + # the flat `ilist` directly shifted every index left, which broke any + # graph whose destination carries leading batch axes (#2382, #2387). + # This mirrors `AdvancedIncSubtensor.perform`. + indices = unflatten_index_variables(ilist, idx_list) + + return mlx_fn(x, indices, y) return advancedincsubtensor diff --git a/tests/link/mlx/test_subtensor.py b/tests/link/mlx/test_subtensor.py index c5bd707c35..b9d814f574 100644 --- a/tests/link/mlx/test_subtensor.py +++ b/tests/link/mlx/test_subtensor.py @@ -347,3 +347,69 @@ def test_mlx_AdvancedIncSubtensor_ignore_duplicates(): assert out.owner.op.ignore_duplicates compare_mlx_and_py([x], [out], [np.zeros(3, dtype=np.float32)]) + + +def test_mlx_AdvancedIncSubtensor_leading_slice(): + """``idx_list`` may place the advanced indices after a leading slice. + + The dispatch used to pass the flat index list straight through, anchoring + the advanced indices at axis 0 and shifting every one of them left, so any + destination carrying leading batch axes failed to broadcast (#2382, #2387). + """ + x = tensor("x", shape=(5, 3, 3), dtype="float32") + idx = pt.arange(3) + out = pt_subtensor.inc_subtensor(x[..., idx, idx], np.float32(1.0)) + + op = out.owner.op + assert isinstance(op, pt_subtensor.AdvancedIncSubtensor) + assert op.idx_list == (slice(None, None, None), 0, 1) + + compare_mlx_and_py([x], [out], [np.zeros((5, 3, 3), dtype="float32")]) + + +def test_mlx_AdvancedIncSubtensor_grad_batched_diagonal(): + """Gradient of a batched diagonal read, the form `specialize` produces. + + `specialize` rewrites ``diagonal(...)`` into advanced indexing, so the + gradient of a log-determinant over a batch of matrices hit #2387 without + the user writing any indexing. + """ + x = tensor("x", shape=(5, 3, 3), dtype="float32") + idx = pt.arange(3) + out = pt.grad(x[..., idx, idx].sum(), x) + + compare_mlx_and_py([x], [out], [np.zeros((5, 3, 3), dtype="float32")]) + + +def test_mlx_AdvancedIncSubtensor_vectorized_scatter(): + """A scatter-add that ``vectorize_graph`` has given a leading batch dim (#2382).""" + from pytensor.graph.replace import vectorize_graph + + n_batch, n_obs, n_levels = 4, 7, 3 + idx = pt.constant( + np.random.default_rng(0).integers(0, n_levels, size=n_obs), dtype="int64" + ) + levels = pt.vector("levels", shape=(n_levels,), dtype="float32") + grad = pt.grad(levels[idx].sum(), levels) + + batched = pt.matrix("batched", shape=(n_batch, n_levels), dtype="float32") + out = vectorize_graph(grad, replace={levels: batched}) + + # Needs the full "MLX" mode: it is `specialize` that collapses the + # `Blockwise` into a single `AdvancedIncSubtensor` with a leading slice. + compare_mlx_and_py( + [batched], + [out], + [np.ones((n_batch, n_levels), dtype="float32")], + mlx_mode="MLX", + ) + + +def test_mlx_AdvancedSetSubtensor_leading_slice(): + """The ``set_instead_of_inc`` branch takes the same ``idx_list`` path.""" + x = tensor("x", shape=(5, 3, 3), dtype="float32") + idx = pt.arange(3) + out = pt_subtensor.set_subtensor(x[..., idx, idx], np.float32(7.0)) + + assert out.owner.op.set_instead_of_inc + compare_mlx_and_py([x], [out], [np.zeros((5, 3, 3), dtype="float32")])