Fixed explicit solves with compiled loops

For long JAX solves, set RunConfig(fixed_loop="scan") to run the fixed explicit mesh through one jax.lax.scan. This avoids tracing a Python copy of every step. It supports Euler, Heun, RK4 and Dormand–Prince, with the same fixed-step updates as the default eager driver. A spatial discretization or nonlocal model can supply the RHS just as an ODE can; the usual stability and discretization requirements still apply.

"""Differentiate a fixed-mesh solve without unrolling its internal steps.

Requires the optional JAX extra. Run ``uv run python examples/fixed_scan.py``.
"""

from __future__ import annotations

import numpy as np


def main() -> None:
    """Check a long fixed solve and its derivative against exponential decay."""
    import jax  # noqa: PLC0415
    import jax.numpy as jnp  # noqa: PLC0415

    from op_engine import CoreSolver, ModelCore  # noqa: PLC0415
    from op_engine.core_solver import RunConfig  # noqa: PLC0415
    from op_engine.model_core import ModelCoreOptions  # noqa: PLC0415

    def final_value(rate: jax.Array) -> jax.Array:
        # Requested outputs and internal step size are static under jit.
        core = ModelCore(
            1,
            1,
            np.asarray([0.0, 10.0]),
            options=ModelCoreOptions(store_history=False),
        )
        core.set_initial_state(jnp.asarray([[1.0]], dtype=rate.dtype))
        CoreSolver(core).run(
            lambda _time, state: -rate * state,
            config=RunConfig(method="rk4", fixed_max_step=0.01, fixed_loop="scan"),
        )
        return core.get_current_state()[0, 0]

    enable_x64 = getattr(jax, "enable_x64", None)
    if enable_x64 is None:
        enable_x64 = jax.experimental.enable_x64
    with enable_x64():
        value, derivative = jax.jit(jax.value_and_grad(final_value))(
            jnp.asarray(0.2, dtype=jnp.float64)
        )
    expected = np.exp(-2.0)
    np.testing.assert_allclose(value, expected, rtol=1e-9)
    np.testing.assert_allclose(derivative, -10.0 * expected, rtol=1e-9)
    print(  # noqa: T201 - command-line example output
        f"Final value: {float(value):.6f}; rate derivative: {float(derivative):.6f}"
    )


if __name__ == "__main__":
    main()

The example checks the solution and parameter derivative against exponential decay. Its 1,000 internal steps do not require 1,000 output records. Requested output times, fixed_max_step, state shape and declared forcing breakpoints are static under JIT. Initial states and RHS parameters can be dynamic arrays. Changing the static mesh can trigger recompilation. Normal JAX restrictions apply to the RHS; Python control flow depending on traced values is unsupported.

fixed_loop="unroll" remains the default. "auto" selects a registered namespace adapter and otherwise retains the eager path. "scan" raises if an adapter is unavailable; JAX is registered lazily. NumPy and CPU PyTorch retain eager execution unless an adapter is explicitly registered. No JAX dependency is added to ordinary core installs.

fixed_checkpoint=True rematerializes scan stages in reverse mode when the adapter supports it. It can reduce stored intermediates at additional runtime cost; the benefit depends on the model. A scan may still materialize internal states before selecting output boundaries, so benchmark memory for large states. ModelCoreOptions(store_history=False) discards the adopted output history.

This is a fixed explicit driver. Adaptive, implicit, IMEX and stochastic solves do not use fixed_loop; unsupported selections raise. For a previously discovered adaptive mesh use adaptive replay. Fixed solves do not create an AdaptiveStepSchedule or adaptive provenance. The FlepiMoP2 provider's prepared execution has its own loop configuration; this option configures CoreSolver.run directly.

Known right-continuous forcing jumps split steps, including jumps between requested outputs, and incoming endpoint stages use the left limit. Scan execution preserves this contract and the fixed_max_step bound. It does not add error control, positivity enforcement or a scientific accuracy guarantee.

Compile-time measurements

The existing replay benchmark also measures live fixed solves in fresh CPU processes, including tracing, compilation, execution and memory metadata:

uv run python scripts/benchmark_replay.py --execution fixed --methods dopri5 \
  --steps 50 100 200 --loops scan unroll --derivatives grad --timeout 120

These are optional measurements, not CI timing thresholds. The regression test checks one scan and a constant traced graph size at 50 and 200 steps. A small linear example can expose driver overhead; its timing is not a performance claim for a larger scientific model.