Strided slice .at(...).add corrupts overlap-add accumulation on MLX 0.31.2

## Summary On `mlx 0.31.2`, repeated `out.at[:, offset:end].add(update)` over strided time slices can corrupt overlap-add accumulation by several thousand times, while equivalent slice assignment and pad-sum implementations reconstruct correctly. This surfaced in `demucs-mlx` as high-amplitude spikes for multi-segment `split=True` inference: https://github.com/ssmall256/demucs-mlx/issues/1 A downstream package release avoids this operation for affected runtimes, but the minimal reproduction below has no dependency on that package and appears reducible to MLX scatter-add behavior. ## Environment Observed on: - `mlx==0.31.2` - Python `3.11.8` - macOS 26 arm64 wheel (`mlx_metal-0.31.2-py3-none-macosx_26_0_arm64`) - Apple Silicon The same reduced overlap-add path did not corrupt output with `mlx==0.31.0` in the same environment and reduced repro. ## Minimal reproduction ```python import numpy as np import mlx.core as mx sr = 44100 segment_length = int(sr * 7.8) overlap = 0.25 stride = int((1 - overlap) * segment_length) length = int(20 * sr) offsets = list(range(0, length, stride)) t = np.arange(length) / sr x_np = 0.5 * np.sin(2 * np.pi * 220 * t) + 0.3 * np.sin(2 * np.pi * 440 * t) x_np = np.stack([x_np, np.roll(x_np, 100)]).astype(np.float32) x_np = (x_np / np.abs(x_np).max() * 0.9).astype(np.float32) x = mx.array(x_np) weight = mx.concatenate([ mx.arange(1, segment_length // 2 + 1), mx.arange(segment_length - segment_length // 2, 0, -1), ], axis=0) weight = (weight / mx.max(weight)).astype(x.dtype) def run_at_add(): out = mx.zeros_like(x) sw = mx.zeros((length,), dtype=x.dtype) for offset in offsets: this_len = min(segment_length, length - offset) end = offset + this_len w = weight[:this_len] update = w.reshape(1, -1) * x[:, offset:end] out = out.at[:, offset:end].add(update) sw = sw.at[offset:end].add(w) mx.eval(out, sw) y = out / sw.reshape(1, -1) mx.eval(y) return np.asarray(y), np.asarray(sw) def run_slice_assign(): out = mx.zeros_like(x) sw = mx.zeros((length,), dtype=x.dtype) for offset in offsets: this_len = min(segment_length, length - offset) end = offset + this_len w = weight[:this_len] update = w.reshape(1, -1) * x[:, offset:end] out[:, offset:end] = out[:, offset:end] + update sw[offset:end] = sw[offset:end] + w mx.eval(out, sw) y = out / sw.reshape(1, -1) mx.eval(y) return np.asarray(y), np.asarray(sw) def run_pad_sum(): out = mx.zeros_like(x) sw = mx.zeros((length,), dtype=x.dtype) for offset in offsets: this_len = min(segment_length, length - offset) end = offset + this_len w = weight[:this_len] update = w.reshape(1, -1) * x[:, offset:end] out = out + mx.pad(update, [(0, 0), (offset, length - end)]) sw = sw + mx.pad(w, [(offset, length - end)]) mx.eval(out, sw) y = out / sw.reshape(1, -1) mx.eval(y) return np.asarray(y), np.asarray(sw) for name, fn in [ ("at_add", run_at_add), ("slice_assign", run_slice_assign), ("pad_sum", run_pad_sum), ]: y, sw = fn() print(name, { "output_peak": float(np.abs(y).max()), "max_abs_error": float(np.max(np.abs(y - x_np))), "sum_weight_min": float(sw.min()), "sum_weight_max": float(sw.max()), }) ``` ## Actual on MLX 0.31.2 ```text at_add {'output_peak': 3349.296143, 'max_abs_error': 3349.27, 'sum_weight_min': 0.000006, 'sum_weight_max': 1.000000} slice_assign {'output_peak': 0.900000, 'max_abs_error': 1.19209e-07, 'sum_weight_min': 0.000006, 'sum_weight_max': 1.000000} pad_sum {'output_peak': 0.900000, 'max_abs_error': 1.19209e-07, 'sum_weight_min': 0.000006, 'sum_weight_max': 1.000000} ``` ## Expected `at_add`, `slice_assign`, and `pad_sum` should all reconstruct the identity input with peak around `0.9` and max absolute error around float32 tolerance. ## Notes - `sum_weight` remains correct in all variants, so the corruption appears specific to the accumulated `out` slice updates. - The failure is reproducible with no model weights or neural-network forward pass. - A downstream workaround is to avoid `.at(...).add(...)` for this accumulation pattern on the affected runtime.