Add MLX backend dispatch for Scan

[cetagostini]

Summary

Adds pytensor/link/mlx/dispatch/scan.py, registering mlx_funcify(Scan) so compiled loops (RNNs/recurrences, autoregressive generation, cumulative scans) can lower to the MLX backend. Previously any Scan raised NotImplementedError: No MLX conversion for the given Op: Scan{...} (JAX and Numba ship a scan dispatch; MLX and PyTorch did not).

It ports the JAX dispatch (link/jax/dispatch/scan.py). MLX has no native general-scan/while primitive to target, so JAX's lax.scan is replaced by a plain Python carry loop that is unrolled into the graph at trace time (since n_steps is statically known); mx.compile then compiles the resulting flat graph. The same trace-recreation semantics are kept: per-step values are stacked, then initial states are prepended / traces truncated to the buffer size.

This is a deliberate trade-off, not a primitive we're missing in PyTensor — it differs from both other backends:

• JAX emits a single lax.scan (one XLA While); the inner graph appears once.
• Numba code-generates a real while i < n_steps loop and JIT-compiles it; again the inner graph appears once and n_steps is a runtime value.
• MLX has no loop op to emit and builds its graph by tracing, so the host-side Python loop unrolls — the inner graph is emitted n_steps times. Consequences: n_steps must be statically known, and graph size / compile time grow linearly with it. (The no-compile path runs the same unrolled ops eagerly, confirming the unrolling happens at trace time rather than inside mx.compile.)

MLX-specific design notes

• Static step count. Under mx.compile, array shapes are concrete but scalar values are not readable. n_steps is therefore taken from a constant n_steps, else inferred from the (full-sized) recurring buffer shapes — which is sound because the MLX linker already excludes scan_reduce_trace_prealloc — else from the sequence length.
• No .at[].set() in MLX, and in-place item assignment aliases buffers under mx.compile, so MIT-MOT writes use a functional scatter-add of the delta (buffer.at[idx].add(vals - buffer[idx])), mirroring the AdvancedIncSubtensor dispatch in the same backend.

Reproducer (now works)

import numpy as np, pytensor, pytensor.tensor as pt
pytensor.config.floatX = "float32"
w = pt.matrix("w"); init = pt.matrix("init")
out, _ = pytensor.scan(lambda p, w: pt.tanh(p @ w),
                       outputs_info=[init], non_sequences=[w], n_steps=3, return_updates=True)
pytensor.function([init, w], out, mode="MLX")  # previously raised NotImplementedError

Known limitation (separate bug, follow-up)

Gradients over sequences reverse the trace, which trips a pre-existing bare-MLX bug: mx.compile miscompiles an elementwise op fed by a negative-stride array, e.g.

import mlx.core as mx, numpy as np
mx.compile(lambda x: 2.0 * x[::-1])(mx.array(np.random.randn(5, 3).astype("float32")))  # tail rows zeroed

This is unrelated to the Scan dispatch (the dispatch logic is correct under the base MLX optimizer query, verified by test_scan_grad_over_sequence). It is captured here as a strict xfail under the full mode="MLX" (test_scan_grad_over_sequence_default_mode) and fixed by a follow-up that materializes negative-stride Subtensor results in the MLX backend.

Test plan

• tests/link/mlx/test_scan.py — SIT-SOT/MIT-SOT (incl. views), sequences + non-seq RNN, multiple/combined recurring states, NIT-SOT-only, dtype preservation, n_steps=0, while-scan raises, gradients; plus the strict xfail above.
• Full MLX suite green (pytest tests/link/mlx/), ruff/pre-commit clean.

🤖 Generated with assistance from Cursor.

[cetagostini]

@ricardoV94 what do you think about this?

[cetagostini]

If this get merge, I'll redo this: ml-explore/mlx#3772