Gate MLX negative-stride materialization before v0.32.0

[cetagostini]

Summary

MLX v0.32.0 fixed compiled-kernel correctness for negative-strided inputs (#3720). This keeps the Subtensor mx.contiguous workaround only for negative-step slices on MLX versions before 0.32.0, so older MLX releases remain correct while fixed releases avoid an unnecessary materialization.

Reproducer and version evidence

import mlx.core as mx
import numpy as np

x = mx.array(np.arange(15, dtype="float32").reshape(5, 3))
out = mx.compile(lambda x: 2.0 * x[::-1])(x)
mx.eval(out)

MLX version	Result
0.31.2	Incorrect: only the first reversed row is correct; the remaining rows are zero.
0.32.0	Correct: exactly matches the NumPy result.

The v0.32.0 release notes list “Fix compiled kernel correctness for negative-strided inputs” (#3720). PyTensor does not declare an MLX minimum and its MLX tests use unversioned pytest.importorskip; more concretely, the current MLX CI job installs mlx<0.29.4 (0.29.3), so it requires the compatibility branch. CI exercises the pre-0.32.0 materialization path; the recorded MLX 0.32.0 test run below exercises the fixed no-copy path.

The regressions run on every MLX version: first- and second-axis negative steps, a non-unit negative step, and simultaneous reversal of both axes, each feeding an elementwise op under mode="MLX"; the previously-xfailed IncSubtensor negative-step gradient path also passes.

Stacked dependency

This branch is stacked on #2258. Until #2258 merges, this PR's diff includes its Scan dispatch and reusable ScanCompatibilityTests coverage; please review #2258 first, then this focused Subtensor change. Ricardo's requested Scan edge cases are present and run with mode="MLX": higher-order derivatives and aliased inner outputs.

Test plan

• Bare mx.compile(lambda x: 2.0 * x[::-1]): RED on MLX 0.31.2 and GREEN on 0.32.0.
• Old-MLX bare checks for first axis, second axis, non-unit step, and both-axis reversal: all fail without materialization and pass with mx.contiguous on MLX 0.31.2.
• PYTENSOR_FLAGS="base_compiledir=/Users/carlostrujillo/Documents/GitHub/_worktrees/pr-2259/.pytensor_compile" /opt/anaconda3/envs/pytensor-pr-2259/bin/pytest -q tests/link/mlx/test_subtensor.py tests/link/mlx/test_scan.py on MLX 0.32.0 — 41 passed, 2 pre-existing xfailed (boolean indexing and inplace operations).
• /opt/anaconda3/envs/pytensor-pr-2259/bin/ruff check and ruff format --check on the touched files.

🤖 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