Materialize negative-stride MLX Subtensor results

[cetagostini]

Summary

mx.compile miscompiles an elementwise op fed by a negative-stride array — trailing entries are zeroed (eager / use_compile=False are correct):

import mlx.core as mx, numpy as np
mx.compile(lambda x: 2.0 * x[::-1])(mx.array(np.arange(15, dtype="float32").reshape(5, 3)))
# tail rows are 0.0 instead of the reversed values

The MLX Subtensor dispatch now materializes reversed slices into a contiguous array (mx.contiguous) so downstream compiled kernels see correct data. Only negative steps are affected (positive steps / basic slices are contiguous and untouched).

Why it matters

• Unblocks Scan gradients over sequences under the full mode="MLX": the backward pass reverses the trace and applies elementwise terms, which previously tripped this bug. The strict xfail added in Add MLX backend dispatch for Scan #2258 (test_scan_grad_over_sequence_default_mode) becomes a normal passing test here.
• Resolves an existing strict xfail (test_mlx_IncSubtensor_negative_step_slice_grad). Its failure — attributed to mx.compile: assigning an elementwise expression to a negative-strided slice writes only one element ml-explore/mlx#3716 (assigning to a negative-strided slice) — was actually this same negative-stride read feeding the elementwise gradient term; the IncSubtensor write to a reversed slice is independently correct (verified). Un-skipped per the "sweep for mis-attributed xfails" convention.

Dependency

Builds on #2258 (MLX Scan dispatch) and should be reviewed/merged after it. The diff will reduce to just the Subtensor change once #2258 lands on main.

Test plan

• New tests/link/mlx/test_subtensor.py::test_mlx_negative_step_slice_elemwise (reversed slice + elementwise under mode="MLX", both axes).
• Un-xfailed test_mlx_IncSubtensor_negative_step_slice_grad and test_scan_grad_over_sequence_default_mode now pass.
• Full MLX suite green (pytest tests/link/mlx/), ruff/pre-commit clean.

🤖 Generated with assistance from Cursor.

Made with Cursor

[ricardoV94]

There are some edge cases of Scan packaged as reusable tests that you can pass mode to, you should try them

[cetagostini]

I'll redo this if mlx patch get merge: ml-explore/mlx#3772