Fix vmap of scatter placing the batch axis wrongly in the updates

[Adityaj0]

Fixes #4321.

Proposed changes

Under vmap the batch dimension becomes 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 was inserted at the front of the source part rather than at src_ax:

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 came out (1, 1, 3), so the scatter wrote with the last two source dimensions transposed and silently produced wrong values.

That is also why it stayed hidden. For axis 0 the intended position and the front coincide. For axis 1 the misplaced singleton lands next to the scattered axis, which is itself size 1, so the shape is unchanged by luck. Only from axis 2 on does the singleton displace a real dimension:

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

The fix inserts the singleton at src_ax within the source part. In the branch where the updates are themselves batched, that insertion shifts their vmap axis, so the index is adjusted before the moveaxis.

Affects every scatter reduction and plain indexed assignment (add, subtract, multiply, maximum, minimum, a[idx] = v). take, take_along_axis and put_along_axis were never affected — they go through Gather / GatherAxis / ScatterAxis rather than the general Scatter.

Before / after on the repro from the issue, where out[0] = 2 * a[0] makes the expected values checkable by hand:

before:  [ 0,  6, 17, 21]
after:   [ 0,  2,  4,  6]     # matches the per-slice loop

Tests

Added test_vmap_scatter_higher_rank to python/tests/test_vmap.py. It sweeps shapes (2,3), (2,3,4) and (2,3,4,5) over every vmap axis and several scatter forms (update derived from the vmapped array, constant update, duplicate indices, and a non-sum reduction), comparing against the equivalent per-slice loop.

I confirmed the test is a real regression test rather than a tautology: with the primitives.cpp change stashed and the extension rebuilt, it fails with

AssertionError: array(False, dtype=bool) is not true : add_derived ax2 (2, 3, 4)

and passes with the change applied.

The existing test_vmap_scatter only covered a 2-D source with in_axes 0 and 1, i.e. exactly the two cases that already worked, which is how the gap survived.

• I have read the CONTRIBUTING document
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (not needed, no API change)
• I have run pre-commit run --all-files to format my code and installed pre-commit prior to committing changes

Verified with a CPU-only build (-DMLX_BUILD_METAL=OFF):

• C++ suite: 247 test cases, 3326 assertions, all pass
• Python: test_vmap, test_ops, test_array, test_autograd, test_einsum, test_linalg, test_blas, test_nn, test_reduce, test_fft (481 tests) pass