MLX: gradient of batched advanced indexing fails in broadcast_shapes

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.

Start by running the reproducer in the issue with the MLX and CVM modes, then trace the AdvancedIncSubtensor gradient through broadcast_shapes. Compare the batched advanced-indexing result with the working diagonal case; done means the MLX gradient evaluates successfully and matches the expected batched result, including the 15.0 sum.