vmap of scatter silently returns wrong values when the vmap axis is >= 2

[Adityaj0]

Describe the bug

mx.vmap of a scatter — a.at[idx].add(...), any other at[] reduction, or a plain a[idx] = v — silently returns wrong values when the vmap axis is 2 or higher. No error is raised.

Scatter::vmap turns the batch dimension into an extra scattered source axis at position src_ax, so the updates need a matching singleton dimension. The updates are laid out as the index dimensions followed by one dimension per source axis, but the singleton is inserted at the front of the source part instead of at src_ax:

// mlx/primitives.cpp, Scatter::vmap
updates = expand_dims(updates, {0, static_cast<int>(inputs[1].ndim())}, stream());

For a source of shape (2, 3, 4) batched on axis 2, the source part of the updates has to be (1, 3, 1) but comes out (1, 1, 3), so the scatter writes with the last two source dimensions transposed.

This is why it only appears from axis 2 on:

• axis 0 — the intended position and the front coincide, so it is correct
• axis 1 — the misplaced singleton lands next to the scattered axis, which is also size 1, so the shape is unchanged by luck
• axis >= 2 — the singleton displaces a real dimension and the result is wrong

To Reproduce

import mlx.core as mx

def unstack(x, axis):
    return [s.squeeze(axis) for s in mx.split(x, x.shape[axis], axis=axis)]

a = mx.arange(2 * 3 * 4, dtype=mx.float32).reshape(2, 3, 4)
f = lambda x: x.at[mx.array([0])].add(x[:1])      # out[0] = 2 * a[0]

out = mx.vmap(f, in_axes=2, out_axes=2)(a)
ref = mx.stack([f(s) for s in unstack(a, 2)], axis=2)

print(out[0, 0, :4])   # array([0, 6, 17, 21], dtype=float32)   <- wrong
print(ref[0, 0, :4])   # array([0, 2, 4, 6], dtype=float32)     <- correct

The expected values are checkable by hand: out[0] = 2 * a[0], and a[0, 0, i] == i, so the first four entries are 0, 2, 4, 6.

Scope, measured on main:

source shape	axis 0	axis 1	axis 2	axis 3
(2,3)	ok	ok	–	–
(2,3,4)	ok	ok	wrong	–
(2,3,4,5)	ok	ok	wrong	wrong

It affects every scatter mode and plain indexed assignment:

a.at[idx].add(u)        # wrong
a.at[idx].subtract(u)   # wrong
a.at[idx].multiply(u)   # wrong
a.at[idx].maximum(u)    # wrong
b[idx] = u              # wrong

Whether the update is derived from the vmapped array or is an independent constant makes no difference.

take, take_along_axis and put_along_axis are not affected — those go through Gather / GatherAxis / ScatterAxis, not the general Scatter primitive.

Expected behavior

mx.vmap(f, in_axes=ax, out_axes=ax)(a) should equal mx.stack([f(s) for s in unstack(a, ax)], axis=ax) for every axis, which is what happens for axes 0 and 1.

Desktop

• OS Version: macOS 26.2
• Version: main @ 8c28c38

Additional context

The existing test_vmap_scatter only exercises a 2-D source with in_axes 0 and 1, which are exactly the two cases that happen to work, so the gap was invisible.

This one is worth flagging as silent rather than loud: nothing raises, the shapes are right, and only the values are wrong, so a batched scatter inside a larger computation will just quietly produce bad numbers.

I have a fix and a regression test ready and will link a PR.

[Adityaj0]

PR: #4322