Honour idx_list in the MLX AdvancedIncSubtensor dispatch - #2397
guillaume-osmo wants to merge 1 commit into
Conversation
, pymc-devs#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) -> pymc-devs#2382, [broadcast_shapes] (4,1) vs (7,3) idx_list=(slice(None), 0, 1) -> pymc-devs#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.
jessegrabowski
left a comment
jessegrabowski
left a comment
There was a problem hiding this comment.
Fix is right, it's the piece of db7fa07 that never reached the mlx dispatch. Two things before merge: slices with bounds still fail because the bounds arrive as mx arrays (pytensor/link/mlx/dispatch/subtensor.py:110), and the four tests are one parametrized test (tests/link/mlx/test_subtensor.py:352). Full mlx suite passes locally.
| return x.at[indices].add(y) | ||
|
|
||
| def advancedincsubtensor(x, y, *ilist, mlx_fn=mlx_fn): | ||
| idx_list = op.idx_list |
There was a problem hiding this comment.
inline this into the signature like incsubtensor above
| 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, |
There was a problem hiding this comment.
drop this comment, it's a changelog. # mirrors AdvancedIncSubtensor.perform is plenty
| # 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) |
There was a problem hiding this comment.
slice bounds come back as 0-d mx arrays here, so x[1:4, idx] still raises. cast the slice components to int (AdvancedSubtensor above has the same bug, factor out a helper and use it in both)
| # 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) |
There was a problem hiding this comment.
use indices_from_subtensor like the rest of the file
| compare_mlx_and_py([x], [out], [np.zeros(3, dtype=np.float32)]) | ||
|
|
||
|
|
||
| def test_mlx_AdvancedIncSubtensor_leading_slice(): |
There was a problem hiding this comment.
collapse the four new tests into one parametrized over inc/set and index form, like test_mlx_AdvancedIncSubtensor1_duplicate_indices below. add x[idx, :] and x[1:4, idx] cases, and a x[1:4, idx] case to test_mlx_AdvancedSubtensor for the read side
| 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")]) |
There was a problem hiding this comment.
zeros with a constant y can't tell set from inc. use arange
| compare_mlx_and_py([x], [out], [np.zeros((5, 3, 3), dtype="float32")]) | ||
|
|
||
|
|
||
| def test_mlx_AdvancedIncSubtensor_grad_batched_diagonal(): |
There was a problem hiding this comment.
this is the same graph as the test above, remove
| compare_mlx_and_py([x], [out], [np.zeros((5, 3, 3), dtype="float32")]) | ||
|
|
||
|
|
||
| def test_mlx_AdvancedIncSubtensor_vectorized_scatter(): |
There was a problem hiding this comment.
remove, x[:, idx] in the parametrized test hits the same idx_list. whether specialize collapses the Blockwise isn't an mlx concern
|
|
||
| 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 |
| [batched], | ||
| [out], | ||
| [np.ones((n_batch, n_levels), dtype="float32")], | ||
| mlx_mode="MLX", |
2 participants
Closes #2382. Closes #2387.
Motivation
Both issues are the same bug.
mlx_funcify_AdvancedIncSubtensorpassed the flat index list straight to the scatter, ignoringop.idx_list.idx_listrecords where the index variables sit among plain slices, so dropping it anchors the advanced indices at axis 0 and shifts every one of them left. Any destination with leading batch axes then fails to broadcast:idx_list(slice(None), 0)[broadcast_shapes] Shapes (4,1) and (7,3) cannot be broadcast(slice(None), 0, 1)[broadcast_shapes] Shapes (5,3) and (3,3) cannot be broadcastBoth are easy to hit without writing any indexing, since
specializerewritesdiagonal(...)into this form — the gradient of a log-determinant over a batch of matrices was enough.Implementation
Call
unflatten_index_variables, mirroringAdvancedIncSubtensor.perform. This is also what the neighbouringAdvancedSubtensordispatch already does (viaindices_from_subtensor).The
set_instead_of_incandignore_duplicatesbranches take the same path, and MLX accepts a slice inside the index tuple for all three (x.at[(slice(None), i, j)].add(y)works).Tests
Four added; all four fail without the change:
test_mlx_AdvancedIncSubtensor_leading_slice— asserts theidx_listshape too, so the test documents the preconditiontest_mlx_AdvancedIncSubtensor_grad_batched_diagonal— the MLX: gradient of batched advanced indexing fails in broadcast_shapes #2387 formtest_mlx_AdvancedIncSubtensor_vectorized_scatter— the MLX AdvancedIncSubtensor fails when Blockwise adds a batch dim #2382 form, under the full"MLX"mode, since it isspecializethat collapses theBlockwisetest_mlx_AdvancedSetSubtensor_leading_slice— theset_instead_of_incbranchDuplicate-index accumulation is unchanged (verified against
np.add.at).Full
tests/link/mlx/suite: no regressions.