Choosing and configuring a solver¶

ModelCore owns output times and state storage. Deterministic systems use CoreSolver and a method selected by RunConfig. Stochastic reaction networks use DirectSSASolver, ThinningSSASolver, TauLeapingSolver, or AdaptiveTauLeapingSolver, because propensities and stoichiometry have different semantics from an ODE right-hand side. In every case the state array selects the Array-API namespace; choosing JAX instead of NumPy does not select a different numerical method.

Deterministic method guide¶

Method	Order	Split/operator input	Typical use
`euler`	1	None	Debugging and first-order reference runs
`heun`	2	None	Default non-stiff explicit integration
`rk4`	4	None	Accurate fixed-step non-stiff integration
`dopri5`	5 (embedded 4)	None	Efficient adaptive non-stiff integration
`imex-euler`	1	One implicit-Euler operator factory	First-order integration of split systems
`imex-heun-tr`	2	One trapezoidal operator factory	Split-step doubling covers explicit and implicit error
`imex-trbdf2`	2	Trapezoidal and BDF2-stage factories	Split systems needing stronger damping
`imex-ark3`	3 (embedded 2)	One implicit-Euler stage factory	Higher-order adaptive split systems
`implicit-euler`	1	Jacobian callable	One-linearization approximation with first-order damping
`trapezoidal`	2	Jacobian callable	One-linearization, second-order integration
`bdf2`	2	Jacobian callable	Uniform, fixed-step, one-linearization integration
`ros2`	2 (embedded 1)	Jacobian callable	L-stable linearly implicit integration
`sdirk2`	2	Full RHS Jacobian and nonlinear solver	L-stable fully nonlinear integration

bdf2 currently requires a uniform output grid and adaptive=False. Its first step uses linearly implicit Euler because no previous state exists yet. The other methods support the built-in adaptive controller, subject to the compiled-control-flow boundary described below.

The linear-multistep design note records the coefficient/history contract, restart rules, and the current no-go decision on exposing BDF3 before variable-step and rejection semantics exist.

Known deterministic forcing jumps¶

For a right-continuous, piecewise-smooth explicit RHS, declare every known jump through RunConfig(forcing_breakpoints=(...)). Euler, Heun, RK4 and Dormand–Prince split fixed/adaptive internal steps at those coordinates. Saved output times remain unchanged; including a jump only in the output grid does not replace this declaration.

import numpy as np
from op_engine import CoreSolver, ModelCore
from op_engine.core_solver import RunConfig

core = ModelCore(1, 1, np.asarray([0.0, 2.0]))
core.set_initial_state(np.ones((1, 1)))
CoreSolver(core).run(
    lambda time, state: np.where(time < 1.0, -state, 0.0 * state),
    config=RunConfig(method="rk4", fixed_max_step=0.125,
                     forcing_breakpoints=(1.0,)),
)
# Final state approximates exp(-1); no extra output row is inserted.

The driver integrates the full duration to each boundary. Incoming RHS stages are capped at its representable left neighbor. Stages starting at a declared jump use a representable right neighbor, ensuring that float32 rounding below a producer's float64 coordinate cannot select the old table value. These adjustments affect RHS evaluation times, not the integrated duration, stored times or state. For declared forcing, first-stage/FSAL cache reuse is disabled throughout the solve, including rejected attempts and replay. Smooth runs with an empty declaration retain their existing kernels and cache behavior.

Times must be finite, strictly increasing and remain distinct/finite in the active float32/float64 state/time precision. Global declarations can extend outside the solve interval. Breakpoints are static metadata, with no state reset, jump impulse, inferred event detection, or gradients with respect to moving jump locations. Callbacks must implement the intended right-continuous regimes and remain smooth enough within each segment for the chosen method.

Functional fixed_explicit_step and adaptive_explicit_step accept the same static declaration. External drivers must clip each step at the next jump; fixed_step_sizes(..., forcing_breakpoints=...) provides a static partition. These kernels adjust stage times but do not shorten the supplied duration. Frozen adaptive replay checks forcing provenance and rejects a mesh that crosses a declared jump; rediscover rather than splitting an accepted replay step. Parameter gradients remain conditional on the frozen mesh and jump times.

The provider merges system.option("forcing_breakpoints") with configured boundaries at each run/preparation, without modifying the reusable engine. This automatically handles op_system time_interpolation: previous tables in flat/PyTree/block fixed runs, prepared execution and adaptive replay. Other deterministic methods and hybrid execution reject declared jumps until their stage/history policies are implemented. Previously the deterministic provider ignored producer declarations, so these unsupported methods now fail clearly. Smooth linear tables publish no jumps and retain their existing behavior.

The implementation follows the known-jump split/cache-refresh principles in Diffrax's controllers and SciML's solver options, with explicit one-sided times for op_system's right-continuous convention.

Stochastic reaction networks¶

All stochastic solvers use one reaction contract: the stoichiometric matrix has shape (n_species, n_reactions), while the propensity function returns the state shape with n_reactions replacing n_species on the configured reaction axis.

Exact direct SSA¶

DirectSSASolver implements Gillespie's direct method. It samples one exponential waiting time and one categorical reaction/batch event, applies exactly one stoichiometric update, and repeats. A draw beyond an output boundary is retained rather than discarded, so adding observation times does not alter the simulated path. Zero total propensity is absorbing when no declared future forcing boundary remains.

Use direct SSA as the reference method for low-copy-number networks, rare-event questions, or checking a tau-leaping approximation. Its cost is proportional to the number of individual events, so tau-leaping is usually preferable when populations and firing rates are high.

import numpy as np

from op_engine import DirectSSASolver, ModelCore, NumpySSASampler

times = np.linspace(0.0, 4.0, 41)
core = ModelCore(n_states=2, n_subgroups=1, time_grid=times)
core.set_initial_state(np.asarray([[20.0], [0.0]]))

# A -> B
stoichiometry = np.asarray([[-1], [1]])


def propensity(_time, state):
    return 0.2 * state[0:1]


DirectSSASolver(core, stoichiometry).run(
    propensity,
    NumpySSASampler(seed=2026),
)

The direct method assumes propensities remain constant in time while the state is unchanged, except at explicitly declared forcing boundaries. Across batch cells it samples from the superposed event process; the result is equivalent to independent direct-method trajectories for independent batch cells.

Piecewise-constant forcing¶

Declare every forcing change through DirectSSAConfig.forcing_breakpoints:

from op_engine import DirectSSAConfig, array_namespace

core = ModelCore(n_states=1, n_subgroups=1, time_grid=np.linspace(0.0, 3.0, 13))
core.set_initial_state(np.zeros((1, 1)))


def forced_birth_rate(time, state):
    xp = array_namespace(state)
    rate = 0.0 if time < 1.0 else (3.0 if time < 2.0 else 1.0)
    return xp.full(state.shape, rate, dtype=state.dtype)


DirectSSASolver(core, np.asarray([[1]])).run(
    forced_birth_rate,
    NumpySSASampler(seed=2026),
    config=DirectSSAConfig(forcing_breakpoints=(1.0, 2.0)),
)

If an event would reach or cross a forcing boundary, the solver advances to that boundary without firing, reevaluates the rates, and redraws. Exponential memorylessness makes this exact for piecewise-constant forcing. Zero-rate intervals advance to the next forcing boundary without drawing randomness.

Forcing is right-continuous: the callback must return the new rates at a breakpoint. In an exact event/breakpoint tie, the forcing boundary wins and the old event is discarded. Observation times remain independent of the forcing schedule, even when an observation coincides with a breakpoint.

Breakpoints must be finite real scalars in strictly increasing order. They are snapshotted as an immutable tuple; a global schedule may include times outside the solve interval. Boundaries at the initial time are already reflected in the first evaluation, and the final boundary causes no extra random draw. An omitted or empty schedule preserves ordinary time-homogeneous execution.

The schedule does not freeze or interpolate callback values. All time dependence between boundaries must be constant for this exactness claim; smoothly varying rates require a different waiting-time method. In particular, op_system defaults to linear interpolation of time-indexed parameters, so listing their grid points as breakpoints does not make direct SSA exact for those interpolated rates. An op_system producer can select time_interpolation: previous to hold table values until the next coordinate. The flepimop2 engine provider consumes their published forcing_breakpoints automatically in pure stochastic mode and combines them with any extra boundaries in the engine configuration. This also supports fixed and adaptive tau-leaping. Any additional direct time dependence must still satisfy the chosen solver's exactness conditions.

Exact SSA for smooth forcing¶

ThinningSSASolver accepts time-dependent propensities and a user-certified total-rate bound. It samples candidates at the bound rate, then evaluates channel rates at the candidate time to reject or apply one event. See the bounded thinning guide for constant and callback bounds, sampling injection, interval expiry, and a runnable smooth-birth example.

Bounded adaptive tau-leaping¶

Both tau-leaping configurations accept the same forcing_breakpoints schedule as direct SSA. Each leap is capped at the next observation or forcing time, as well as the method's usual step-size limit. The next leap evaluates the right-continuous rates at that boundary. Boundaries count toward the fixed or adaptive method's max_steps guard for the entire output interval; they do not reset it.

For example, the piecewise birth callback above can also use fixed tau-leaping:

from op_engine import NumpyPoissonSampler, TauLeapingConfig, TauLeapingSolver

core = ModelCore(n_states=1, n_subgroups=1, time_grid=np.linspace(0.0, 3.0, 13))
core.set_initial_state(np.zeros((1, 1)))

TauLeapingSolver(core, np.asarray([[1]])).run(
    forced_birth_rate,
    NumpyPoissonSampler(seed=2026),
    config=TauLeapingConfig(max_step=0.1, forcing_breakpoints=(1.0, 2.0)),
)

Adaptive tau-leaping uses these caps for both accepted leaps and rejected proposals. A critical event tied with a forcing boundary is discarded, while noncritical Poisson counts cover the segment ending there. During exact fallback, a pending event survives observation times but is discarded at or beyond the next forcing change. Zero total propensity waits for that change without drawing randomness; without a future boundary it remains absorbing.

Boundary caps keep a leap from using the old rate across a known change. The usual tau-leaping approximation still freezes state-dependent propensities within each leap. Smooth time variation between boundaries also remains an approximation for leaps and is unsupported by adaptive exact fallback. For that fallback, use rates constant in time between declared forcing changes while the state is unchanged.

AdaptiveTauLeapingSolver implements the species-based pre-leap selector of Cao, Gillespie, and Petzold together with their non-negative modified Poisson algorithm. It needs both net stoichiometry and non-negative reactant stoichiometry. They are intentionally separate: net changes cannot reveal catalytic reactants or how many source molecules a reaction consumes.

from op_engine import (
    AdaptiveTauLeapingConfig,
    AdaptiveTauLeapingSolver,
    NumpyPoissonSampler,
    NumpySSASampler,
)

# Net A -> B changes and consumed reactants, respectively.
stoichiometry = np.asarray([[-1], [1]])
reactants = np.asarray([[1], [0]])

solver = AdaptiveTauLeapingSolver(core, stoichiometry, reactants)
solver.run(
    propensity,
    NumpyPoissonSampler(seed=2026),
    NumpySSASampler(seed=2027),
    config=AdaptiveTauLeapingConfig(leap_tolerance=0.03),
)

The non-negativity strategy has three explicit layers:

Pre-leap selection. Species drift and variance from noncritical reactions bound expected and stochastic population changes. The leap_tolerance controls this approximation.
Bounded critical events. A channel is critical when fewer than critical_threshold firings could exhaust a reactant. Across all critical channel/batch events, at most one is selected by the exact-event sampler in a leap. If all active channels are critical, the method is direct SSA.
Post-leap rejection. If noncritical Poisson firings still propose a negative population, the proposal is discarded, the attempted tau is halved, and fresh samples use new deterministic draw indices. No population or firing count is clipped.

When the selected tau is cheaper to execute event by event, exact_fallback_multiplier switches temporarily to direct SSA. Tightening leap_tolerance therefore approaches, and eventually crosses into, the exact reference path on low-count networks.

Per-channel binomial firing was considered but is not used here. Although a binomial draw bounds one channel, independent binomials can still collectively overconsume a reactant shared by several channels. A true coupled-binomial scheme would need a distinct sampling contract. The critical-event bound plus explicit post-leap rejection gives a documented non-negative method without silently changing those reaction dependencies.

Fixed-step tau-leaping¶

Fixed-step explicit tau-leaping approximates the number of firings in each reaction channel over a time interval with an independent Poisson draw. Supply a stoichiometric matrix with shape (n_species, n_reactions) and a propensity function whose reaction axis has length n_reactions:

import numpy as np

from op_engine import (
    ModelCore,
    NumpyPoissonSampler,
    TauLeapingConfig,
    TauLeapingSolver,
)

times = np.linspace(0.0, 4.0, 41)
core = ModelCore(n_states=2, n_subgroups=1, time_grid=times)
core.set_initial_state(np.asarray([[1_000.0], [0.0]]))

# A -> B
stoichiometry = np.asarray([[-1], [1]])


def propensity(_time, state):
    return 0.2 * state[0:1]


solver = TauLeapingSolver(core, stoichiometry)
solver.run(
    propensity,
    NumpyPoissonSampler(seed=2026),
    config=TauLeapingConfig(max_step=0.05),
)

max_step=None takes one leap per output interval. Setting max_step divides each interval into smaller leaps and clips the final leap so that it lands on the requested output time. The propensity is evaluated at the beginning of each leap.

The core solver does not own random-number state. Instead, it receives a PoissonSampler, keeping PRNG details outside the numerical method. A JAX run can use an explicit key and the stable leap index:

key = jax.random.key(2026)


def poisson(mean, step_index):
    return jax.random.poisson(jax.random.fold_in(key, step_index), mean)

This first implementation deliberately rejects a leap that produces a negative population. It neither clips counts nor changes the process silently. Reduce max_step for a better fixed-tau approximation, or use AdaptiveTauLeapingSolver when automatic leap selection and critical-event handling are required.

Tau-leaping is currently an eager execution path: validation reads sampled counts and proposed populations back to Python. JAX arrays stay in their native namespace, but a sampled stochastic trajectory does not provide an ordinary pathwise jax.grad derivative. Differentiable deterministic counterparts can continue to use CoreSolver; gradient estimators for stochastic paths belong at a higher inference/provider layer.

Explicit deterministic methods¶

Euler, Heun, classic RK4, and Dormand--Prince need only an RHS. The default is fixed-step Heun, with one step per output interval:

import numpy as np

from op_engine import CoreSolver, ModelCore
from op_engine.core_solver import AdaptiveConfig, RunConfig

times = np.linspace(0.0, 4.0, 41)
core = ModelCore(n_states=1, n_subgroups=1, time_grid=times)
core.set_initial_state(np.asarray([[1.0]]))


def rhs(_time, state):
    return -0.5 * state


config = RunConfig(
    method="heun",
    adaptive=True,
    adaptive_cfg=AdaptiveConfig(rtol=1e-6, atol=1e-9),
)
CoreSolver(core).run(rhs, config=config)

With adaptive=False, every adjacent pair of output times defines one solver step by default. Set fixed_max_step to bound explicit Euler, Heun, RK4, or Dormand--Prince integration steps independently of the output grid:

config = RunConfig(method="rk4", fixed_max_step=0.25)

Each output interval then uses as many full fixed_max_step steps as fit plus a final remainder step that lands exactly on the next output time. Only the requested output states are stored. Output grids aligned with the fixed-step mesh can therefore be made more or less frequent without changing states at shared times. fixed_max_step is incompatible with adaptive=True; adaptive methods use their controller configuration instead.

With adaptive=True, the controller may take multiple internal steps but still stores only the requested output times.

Heun, RK4, and Dormand--Prince share one validated explicit Runge--Kutta tableau kernel. rk4 takes four RHS stages per fixed step. Under adaptivity it uses step doubling to estimate error, so dopri5 is normally the more efficient adaptive choice: its embedded fourth-order formula estimates error using the same seven stages as its fifth-order solution. Dormand--Prince also reuses its final stage as the first stage of the next accepted step (FSAL), including across output-time boundaries.

dopri5, rk45, and dormand-prince select the same canonical method. Dense output is not currently exposed; the solver lands on and stores the requested output times instead of interpolating between accepted internal steps.

IMEX methods and operator factories¶

IMEX methods solve a split system

[ y' = A(t, y)y + F(t, y), ]

where the RHS callable supplies F and an operator factory supplies the implicit part. A resolved operator tuple is either (L, R) or (predictor, L, R), with the core solving L @ y_next = R @ x along the configured operator axis.

Use a StageOperatorFactory whenever the output spacing is non-uniform or adaptive stepping is enabled. The factory receives the attempted full-step size, a stage scale, and StageOperatorContext, so it can rebuild operators for the actual stage time and state.

import numpy as np

from op_engine import CoreSolver, ModelCore, build_diffusion_matrix
from op_engine.core_solver import OperatorSpecs, RunConfig
from op_engine.matrix_ops import (
    StageOperatorContext,
    make_stage_operator_factory,
)

n_cells = 32
dx = 1.0 / n_cells
diffusivity = 0.02
times = np.linspace(0.0, 1.0, 21)

core = ModelCore(n_states=n_cells, n_subgroups=1, time_grid=times)
core.set_initial_state(np.ones((n_cells, 1)))


def reaction(_time, state):
    return -0.1 * state


def diffusion(ctx: StageOperatorContext):
    return build_diffusion_matrix(
        n_cells,
        dx,
        diffusivity,
        bc="neumann",
        reference=ctx.y,
    )


trapezoidal = make_stage_operator_factory(diffusion, scheme="trapezoidal")
config = RunConfig(
    method="imex-heun-tr",
    operators=OperatorSpecs(default=trapezoidal),
)
CoreSolver(core, operator_axis="state").run(reaction, config=config)

Adaptive Heun--trapezoidal error control¶

Set adaptive=True and supply a timestep-dependent trapezoidal factory for imex-heun-tr. Its estimator compares one complete split step full with two complete half steps refined, including their explicit and implicit stages. Each factory call receives the actual timestep, stage time and candidate state.

The accepted state remains full, with local error estimated by (4/3) * (refined - full). For a smooth second-order problem the refinement difference has cubic leading order; the timestep controller uses exponent -1/3. The factor 4/3 estimates the full-step error. A factor 1/3 would estimate the error of the two-half-step state instead.

This replaces the old explicit-only low/high difference, which could vanish when the explicit RHS was zero even though the implicit time-discretization error was nonzero. Adaptive meshes and results can therefore change, including for mixed systems. Fixed stepping and replay keep the original full-step formula and perform two RHS evaluations and two implicit solves per step. Live adaptive estimation performs six of each per attempted step, including rejected attempts. Static operator pairs remain supported for uniform fixed steps; adaptive runs require a factory so half-step operators are rebuilt.

The required CPU qualification checks pure implicit and mixed scalar decay against exact exponentials, cubic estimator scaling, tolerance refinement, factory stage contexts and replay equivalence. Structured transport checks provide independent matrix-exponential references. These are local truncation error estimates, rather than global accuracy bounds. Trapezoidal integration also retains its usual lack of L-stability. State-dependent or nonsmooth assembly requires an independent convergence check for the actual problem; metadata or a successful controller decision does not qualify it.

Step doubling also supplies implicit error estimates in GSL. An embedded additive pair can reduce estimation cost, as described by ARKODE, but replacing the Heun--trapezoidal stages would change this method's fixed and replay behavior. The existing imex-ark3 provides a separately qualified additive-tableau method. Mesh discovery provenance records the Heun--trapezoidal estimator and the partitions it covers; see replay provenance.

For a non-uniform axis, pass strictly increasing cell-center coordinates and set dx=None:

centers = np.asarray([0.0, 0.08, 0.21, 0.5, 1.0])
operator = build_diffusion_matrix(
    centers.size,
    None,
    diffusivity,
    grid=centers,
    bc="neumann",
)

The no-flux operator is conservative and self-adjoint under the inferred cell-volume weights. It is generally not symmetric in the unweighted Euclidean inner product. Periodic non-uniform grids need an explicit wrap spacing, which cannot be inferred from center coordinates alone, and are rejected.

Use generate_adaptive_grid to create a curvature-weighted point grid from a vectorized profile callable:

from op_engine import generate_adaptive_grid

centers = generate_adaptive_grid(
    lambda x: np.exp(-10.0 * x * x),
    (-1.0, 1.0),
    65,
    smoothing_window=5,
    minimum_spacing=1e-3,
)

The generator returns exactly the requested number of NumPy coordinates, including both interval endpoints. It is eager preprocessing rather than part of the differentiable solve; the resulting static geometry can drive NumPy, JAX, Torch, or another supported Array-API backend. The sampled-data companion is generate_adaptive_grid_from_data.

For imex-trbdf2, provide both stage schemes:

implicit_euler = make_stage_operator_factory(diffusion, scheme="implicit-euler")
config = RunConfig(
    method="imex-trbdf2",
    operators=OperatorSpecs(tr=trapezoidal, bdf2=implicit_euler),
)

imex-ark3 uses the five-stage ARS(4,4,3) additive Runge--Kutta tableau from Ascher, Ruuth, and Spiteri (1997). Its implicit half is L-stable and stiffly accurate. A shared c=1/2 stage supplies a second-order explicit/implicit midpoint estimate, so adaptive error control does not use step doubling.

implicit_euler = make_stage_operator_factory(diffusion, scheme="implicit-euler")
config = RunConfig(
    method="imex-ark3",
    operators=OperatorSpecs(default=implicit_euler),
)

This method always requires a StageOperatorFactory, including on a uniform fixed-step grid. For each non-explicit stage the factory receives the attempted full-step size, the DIRK diagonal coefficient as scale, the actual stage time, and the assembled stage base in StageOperatorContext.y. It must build an implicit-Euler left operator I - dt * scale * A. Predictor tuples are not accepted, and the returned right operator is not applied: the additive kernel has already assembled all previous explicit and implicit stage contributions. The implicit derivative is then recovered from the solved stage equation, which avoids a duplicate callback for A @ y and keeps the method portable across Array-API namespaces.

imex-ars443 and ars443 are aliases for imex-ark3.

The flepimop2 provider can compile typed op_system axis_kernel, jump_integral, advection, and diffusion descriptors into these same factories. Descriptor parsing, axis-type/coordinate resolution, and layout lifting belong to the provider; portable jump assembly follows op_system's public reference semantics, while stepping methods remain in op_engine.

Linearly implicit methods¶

The methods historically named implicit-euler, trapezoidal, and bdf2, as well as the Rosenbrock method, take a Jacobian callable. The callable returns an operator acting along operator_axis:

from op_engine import array_namespace


def jacobian(_time, state):
    xp = array_namespace(state)
    return xp.multiply(xp.eye(state.shape[0], dtype=state.dtype), -0.5)


config = RunConfig(method="ros2", jacobian=jacobian)
CoreSolver(core, operator_axis="state").run(rhs, config=config)

ros2 uses two increment stages and one Jacobian evaluation per attempted step. Both stages solve with the same matrix I - (1 - 1/sqrt(2)) * dt * J; the accepted second-order state and embedded first-order state provide the adaptive error estimate. The coefficient table, formal orders, and controller order are declared together and validated when the module is imported.

Dense Jacobians stay in the state's namespace. SciPy sparse operators are an optional NumPy acceleration path and are not a JAX differentiation path.

The first three methods perform a single linearization; they do not iterate a nonlinear residual to convergence. sdirk2 instead solves each stage through the nonlinear solver contract, with a distinct full flattened Jacobian and explicit convergence diagnostics.

Adaptivity and Array-API backends¶

AdaptiveConfig sets tolerances and attempt limits; DtControllerConfig sets step-size bounds and growth factors. Their static scalar values are validated when the configuration is constructed. Checks that depend on the time grid, state shape, or resolved operators occur when CoreSolver builds its run plan.

Fixed-step explicit, dense IMEX, and dense linearly implicit methods use ordinary Array-API operations. With JAX arrays, those same methods can run under jax.jit and be differentiated with jax.grad; no JAX-specific solver implementation is selected.

The built-in adaptive controller uses eager Python loops and scalar acceptance decisions. It preserves JAX arrays, and an eager jax.grad differentiates the accepted sequence of native steps, but the live controller is not itself JIT-compatible.

For compiled differentiation, record the accepted mesh at representative parameters and replay it:

solver = CoreSolver(reference_core)
solver.run(rhs, config=config)
schedule = solver.last_adaptive_schedule
assert schedule is not None

# Construct a fresh core and solver inside the function being transformed.
solver = CoreSolver(differentiable_core)
diagnostics = solver.replay_adaptive_schedule(rhs, schedule, config=config)

jax.jit(jax.value_and_grad(...)) can trace the replay because its step count and step sizes are static, while array-valued model inputs remain dynamic. The gradient is conditional on that mesh: replay does not differentiate the accept/reject decisions. Refresh the schedule when parameters, tolerances, the method, or other model structure change materially.

For SDIRK2, replay returns array-valued nonlinear diagnostics as part of the compiled result. Inspect diagnostics.require_converged() after execution and discard the replayed state and derivatives if any frozen-mesh stage failed.

A live compiled controller or specialized adjoint implementation can still belong in an external provider and may return a trajectory through ModelCore.apply_trajectory. It should reuse the portable method semantics rather than define a JAX-only numerical method.

See Backend and solve-strategy boundaries for the complete portability contract and the biogeochemical tutorial for a larger IMEX splitting example.