[BUG] Silently wrong grad and vmap for a strided slice that selects exactly one element

[mentria-ai]

☑️ I understand it is strictly prohibited to use AI to write issues.

Describe the bug
A strided slice that selects exactly one element out of a span of two or more
stores a (start, end, strides) triple that no longer round-trips, and every
consumer that re-derives the region from it is wrong. Nothing raises. Three
symptoms, one cause:

1. mx.grad returns a silently wrong gradient. The forward value is correct,
   so this is invisible until a downstream number is quietly a few percent out.
   The cotangent is broadcast over the whole half-open span [start, stop)
   instead of being deposited at start.
2. mx.vmap returns a wrong value AND a wrong shape. Not gradient-only.
3. Negative strides silently zero the gradient rather than misplacing it.

Reproduces on both the CPU and the GPU stream, on float32 / float16 /
bfloat16, and on any axis.

To Reproduce

1. Wrong gradient.

import mlx.core as mx

x  = mx.zeros((2, 3))
Wa = mx.array([[  1.,   2.,   3.]])
Wb = mx.array([[100., 101., 102.]])

print(mx.grad(lambda z: (z[0::2] * Wa).sum() + (z[1::2] * Wb).sum())(x))
# got      [[1., 2., 3.], [101., 103., 105.]]     <- row 1 is Wa + Wb
# expected [[1., 2., 3.], [100., 101., 102.]]

z[0::2] selects row 0 only, so Wa belongs in row 0; it is written to rows 0
and 1, and Wb is then added on top of row 1.

The broadcast-over-the-span nature is clearer with a longer span:

x = mx.zeros((8, 2)); W = mx.array([[1., 2.]])
print(mx.grad(lambda z: (z[2:5:3] * W).sum())(x))
# rows 2, 3 AND 4 come back [1., 2.];  expected row 2 only
# control: z[2:3:3] (span 1) correctly writes row 2 alone

2. Wrong vmap value and shape — this one is not gradient-only:

x = mx.arange(24, dtype=mx.float32).reshape(4, 2, 3)
print(x[0][0::2].shape)                      # (1, 3)  - the un-vmapped slice
print(mx.vmap(lambda t: t[0::2])(x).shape)   # (4, 2, 3)   expected (4, 1, 3)
# the values are both rows, not row 0; nothing raises

y = mx.arange(48, dtype=mx.float32).reshape(4, 4, 3)   # control: 2 elements out
print(mx.vmap(lambda t: t[0::2])(y).shape)   # (4, 2, 3)   correct

3. Negative stride, gradient silently zero:

x = mx.zeros((2, 2)); W = mx.array([[1., 2.]])
print(mx.grad(lambda z: (z[::-2] * W).sum())(x))
# got      [[0., 0.], [0., 0.]]
# expected [[0., 0.], [1., 2.]]

Which combinations are affected

For [start:stop:stride] on an axis of length n, span = stop - start
(stop defaulting to n):

• positive stride — wrong iff ceil(span / stride) == 1 and span >= 2;
  for an open-ended [start::stride] that is 2 <= n - start <= stride.
• negative stride — wrong iff the slice selects exactly one element.

The vmap symptom has the same trigger condition.

Stride 1 is never affected; a slice selecting two or more elements is never
affected. Exhaustively verified over lengths 2..8 x strides 1..4 x
starts 0..stride-1 x both axes of a 2-D array: 40 of 140 combinations have a
wrong gradient, all of them satisfying the rule above, identical on CPU and GPU.

stride 2 (parenthesised: number of elements selected)

length	start 0	start 1
2	WRONG (1)	ok (1)
3	ok (2)	WRONG (1)
4	ok (2)	ok (2)
5	ok (3)	ok (2)

stride 4

length	start 0	start 1	start 2	start 3
2	WRONG (1)	ok (1)	ok (0)	ok (0)
3	WRONG (1)	WRONG (1)	ok (1)	ok (0)
4	WRONG (1)	WRONG (1)	WRONG (1)	ok (1)
5	ok (2)	WRONG (1)	WRONG (1)	WRONG (1)
6	ok (2)	ok (2)	WRONG (1)	WRONG (1)
7	ok (2)	ok (2)	ok (2)	WRONG (1)
8	ok (2)	ok (2)	ok (2)	ok (2)

Likely cause

normalize_slice in mlx/ops.cpp collapses the stride when an axis yields a
single element, but leaves the bound alone (and takes stop by value, so it
could not write a corrected bound back anyway):

// mlx/ops.cpp
normalize_slice(const Shape& shape, Shape& start, Shape stop, Shape& strides) {
  ...
    // Simplify the stride if it's unused
    if (out_shape[i] == 1) {
      strides[i] = 1;
    }

The Slice primitive then stores a (start, end, strides) triple that no longer
describes the same element set — for z[0::2] on a length-2 axis it stores
start=[0,0], end=[2,3], strides=[1,1], which names two rows. The forward is
unaffected because the correct out_shape is carried separately as the array's
shape. But anything that re-derives the region from the stored triple gets the
oversized one, and there are two such consumers:

Slice::vjp (mlx/primitives.cpp) passes the triple to slice_update:

auto out = zeros_like(primals[0], stream());
return {slice_update(
    out, cotangents[0], start_indices_, end_indices_, strides_, stream())};

slice_update re-normalizes it to upd_shape = [2, 3], broadcasts the [1, 3]
cotangent up to it, and — since upd_shape == src.shape() — returns it whole.
Hence symptom 1.

Slice::vmap (mlx/primitives.cpp) passes the triple back into slice():

auto start = start_indices_;
auto stop = end_indices_;
auto strides = strides_;
...
return {{slice(input, start, stop, strides, stream())}, {ax}};

which re-normalizes to out_shape = [2, 3] and returns two rows. Hence symptom 2.

For negative strides the same collapse writes +1 in place of -k, losing the
sign; re-normalization then yields an empty region — symptom 3.

A fix presumably wants to make the simplification shape-preserving: take stop
by reference and, in the same branch, set stop[i] = start[i] + 1 (and
start[i] - 1 with strides[i] = -1 for the negative case) — or drop the
simplification entirely. Either fixes all three symptoms without touching a
backend kernel.

Why CI does not catch it

test_slice_grads (python/tests/test_autograd.py) looks like the only strided-
slice gradient test. It uses a[5:-6:-1] on a length-5 array (5 elements out)
and a[4:-5:-2] on a length-4 array (2 elements out) — neither ever produces
out_shape == 1 with |stride| > 1, so the branch above is never exercised
under autodiff.

For what it is worth, git log suggests the VJP symptom arrived with #1727
("Simplify + speedup slice::vjp by using slice_update instead of scatter",
merged 2024-12-23); the scatter implementation it replaced built an explicit
index list and so did not re-derive the region.

Expected behavior
The gradient should place each cotangent element at the position the forward
slice read it from, and vmap of a slice should return the batched version of
what the un-batched slice returns — for any (start, stop, stride).

Desktop (please complete the following information):

• OS Version: macOS 26.5.1 (Apple M4 Pro)
• Version: mlx 0.32.1 (pip). normalize_slice, slice, slice_update,
  Slice::vjp and Slice::vmap are byte-identical between the v0.32.1 and
  v0.32.2 tags and main at 052e77db, so this is almost certainly still
  present on main; I diffed the source rather than building it.

Additional context
Found in a chunkwise DeltaNet training kernel: a block-recursive matrix inverse
splits blocks even/odd with inv[:, 0::2] / inv[:, 1::2], and at the last level
of the recursion that axis has length 2. Every earlier level was correct, so the
forward and most gradients were exact while three of the input gradients drifted
0.5-2.5 % — indistinguishable from ordinary numerical fragility. Swapping the
strided slice for a reshape split (or mx.take) fixes it and changes no
forward value. The vmap symptom turned up afterwards, while reading
Slice::vmap to understand the VJP one.

Workarounds that all give the correct gradient today, if anyone hits this before
a fix lands: mx.take, a reshape split, fancy indexing, mx.split, or writing
the slice with an explicit unit stride (z[i:i+1]).