MLX: gradient of batched advanced indexing fails in broadcast_shapes

[jessegrabowski]

The gradient of advanced integer indexing drops the leading batch dimensions on MLX, so the AdvancedIncSubtensor scattering the adjoint back tries to broadcast a (5, 3) against the (3, 3) core and dies. Easy to hit without writing any indexing yourself: specialize rewrites diagonal(cholesky(X)) into this form, so the gradient of a log-determinant over a batch of matrices fails.

import numpy as np
import pytensor
import pytensor.tensor as pt

X = pt.tensor("X", shape=(5, 3, 3), dtype="float32")
idx = pt.arange(3)
g = pt.grad(X[..., idx, idx].sum(), X)
Xv = np.zeros((5, 3, 3), dtype="float32")

print(pytensor.function([X], g, mode="CVM")(Xv).sum())  # 15.0
print(pytensor.function([X], g, mode="MLX")(Xv).sum())
# ValueError: [broadcast_shapes] Shapes (5,3) and (3,3) cannot be broadcast.

pt.diagonal(X, axis1=-2, axis2=-1) and its gradient are both fine; it's the advanced-indexing spelling the rewrite produces that breaks. Writing the diagonal as (X * pt.eye(3)).sum(-1) avoids it.

[jessegrabowski]

Root cause, and it isn't the batch dimensions — my original diagnosis above is wrong. mlx_funcify_AdvancedIncSubtensor hands the scatter only the array-valued indices, but slices and plain integers live in op.idx_list and have to be spliced back in. AdvancedSubtensor, right above it in the same file, does exactly that with indices_from_subtensor; this one doesn't, so x[:, idx] = y becomes x[idx] = y and writes along the leading axis.

Nothing to do with gradients, and no batch dimension needed:

import numpy as np
import pytensor
import pytensor.tensor as pt

X = pt.matrix("X", shape=(6, 6), dtype="float32")
y = pt.matrix("y", shape=(6, 3), dtype="float32")
Xv, yv = np.zeros((6, 6), dtype="float32"), np.ones((6, 3), dtype="float32")
out = pt.set_subtensor(X[:, np.array([0, 1, 2])], y)  # op.idx_list == (slice(None), 0)

print(pytensor.function([X, y], out, mode="CVM")(Xv, yv).sum())  # 18.0
print(pytensor.function([X, y], out, mode="MLX")(Xv, yv).sum())
# ValueError: [broadcast_shapes] Shapes (6,3) and (3,6) cannot be broadcast.

The gradient MWE in the issue reaches it because the adjoint scatters through a slice. So the title is misleading too — it's any AdvancedIncSubtensor whose index tuple isn't purely array-valued.

Potential fix (requires testing):

def advanced_inc_subtensor(x, y, *ilist):
    indices = indices_from_subtensor(ilist, op.idx_list)
    if not isinstance(indices, tuple):
        indices = (indices,)
    ...  # existing set / ignore_duplicates / scatter-add branches

With that splice, X[:, idx] = y, X[idx, :] = y.T, X[:, idx] += y and the gradient MWE above all match the CPU result.