ENHANCEMENT: Scan has no MLX dispatcher — stateful recurrent logps fail to compile

Part of #2085.

Current state

pytensor/link/mlx/dispatch/ registers no handler for the Scan Op. Vectorisable scan patterns (e.g. AR(1) with taps=[0, -1]) get rewritten away to Subtensor + Composite before reaching the backend, so they "work". But genuinely stateful scans — where each step's output feeds the next via outputs_info — survive the rewriter and crash at compile time.

Reproducer

import pytensor, pytensor.tensor as pt

T = 100
y = pt.vector("y", dtype="float32")
rho = pt.scalar("rho", dtype="float32")

def step(y_t, h_prev, rho):
    return pt.tanh(rho * h_prev + y_t)

h_init = pt.zeros((), dtype="float32")
hs, _ = pytensor.scan(
    fn=step, sequences=[y], outputs_info=[h_init], non_sequences=[rho],
)
total = hs.sum()

pytensor.function([y, rho], total, mode="MLX")
# NotImplementedError: No MLX conversion for the given Op:
# Scan{scan_fn, while_loop=False, inplace=none}.

Why this matters for PyMC

A lot of PyMC distributions and likelihoods build their logp with pytensor.scan:

• AR(p), MA(q), ARMA models with non-trivial dependence
• HMM forward / backward
• LSTM / GRU likelihoods
• Kalman filter (linear-Gaussian state space)
• Particle filter / SMC step
• Any user-written recurrent likelihood

All of these are unusable on the MLX backend today.

Proposed change

There are two reasonable directions; either is acceptable, both could ship.

Option A — Native MLX Scan dispatcher (Python loop with MLX kernels)

The simplest implementation: dispatch each step into MLX, run the loop in Python, accumulate outputs. This is what the JAX backend used to do before lowering Scan to lax.scan. MLX's lazy compute helps here because the inner kernel launches queue up and the host-side loop is mostly bookkeeping.

Sketch:

# pytensor/link/mlx/dispatch/scan.py
from pytensor.scan.op import Scan

@mlx_funcify.register(Scan)
def mlx_funcify_Scan(op, node, **kwargs):
    inner_fn = mlx_funcify(op.fgraph)
    n_steps_idx = ...  # from op.info
    n_seqs = op.info.n_seqs
    n_mit_mot = op.info.n_mit_mot
    n_mit_sot = op.info.n_mit_sot
    n_sit_sot = op.info.n_sit_sot
    n_nit_sot = op.info.n_nit_sot

    def scan(*inputs):
        # unpack n_steps, sequences, outputs_info, non_sequences
        # allocate output buffers (mx.zeros)
        # loop in Python; each step runs the inner_fn
        # carry state through outputs_info slots
        # return final stacked outputs
        ...
    return scan

The harder bits are the multi-tap (mit_sot) and nit_sot slot bookkeeping. The Numba and JAX backends both have this code — the MLX dispatcher can crib heavily from pytensor/link/jax/dispatch/scan.py since the semantics are the same. The only platform-specific bit is that MLX arrays use functional updates (mx.array.at-style or full-array reassignment), not in-place arr[i] = ….

Option B — Lower Scan to a graph rewrite when possible

Many scans the user writes can be rewritten to vectorised tensor ops (pytensor.scan → Subtensor + Composite, like the AR(1) case already handles). PyTensor already has some of this machinery; extending the rewriter to cover more cases would shrink the set of scans that hit the backend. This is complementary to Option A but doesn't eliminate the need for one.

Option C (interim) — better error message

Until either A or B lands, replace the generic NotImplementedError with a backend-aware error that tells the user which Op is missing and points them at this issue / the JAX backend as a workaround:

def mlx_funcify(op, node=None, **kwargs):
    raise NotImplementedError(
        f"The MLX backend has no implementation for Op `{type(op).__name__}`. "
        "Stateful Scan is tracked in #<this-issue>; for now use mode='JAX' or "
        "mode='NUMBA' for graphs containing it."
    )

Acceptance criteria

• Repro above compiles and produces results within fp32 tolerance of the Numba backend.
• Coverage tests for: single-output sit_sot, multiple outputs_info, non_sequences, n_steps runtime input, scan over multiple sequences with different taps.
• Numerics tests against Numba reference for a small AR(1) with outputs_info and a 2-state HMM forward pass.

Start by reading pytensor/link/jax/dispatch/scan.py and the MLX dispatcher under pytensor/link/mlx/dispatch/, then run the reproducer in the issue. Implement Scan support or the specified interim handling, and use the acceptance criteria to verify compilation, fp32 agreement with Numba, and coverage for sit_sot, outputs_info, non_sequences, runtime n_steps, multiple sequences, and a two-state HMM.