Differentiable training

Train dynamics (or observation operators) through the filter: run a full assimilation over T windows, score it with a log-likelihood or truth-matching loss, and differentiate with respect to the model parameters. Two gradient strategies are provided:

• differentiable_assimilate — a fixed-shape jax.lax.scan over the forecast–analysis cycle with the full forward+backward tape (optionally jax.checkpoint-ed for O(√T) memory). Use when you want exact gradients, including the cross-time terms that flow through the analysis updates.
• road_enkf_loss_and_grad / road_enkf_grad_step — the ROAD-EnKF local-gradient strategy, with jax.lax.stop_gradient between cycles. Backward-pass memory is O(Nₑ · Nₓ) independent of T, at the cost of dropping cross-time gradient terms.

Both refuse stochastic filters (StochasticEnKF, ETKF_Livings) and the stochastic AdditiveInflator at call time — randomness in the analysis breaks the deterministic tape these strategies rely on. The deterministic ETKF analysis core (rank-N_y QR + symmetric eigh) is what keeps the exact-gradient path differentiability-safe.

Full-tape assimilation

filterax — differentiable ensemble Kalman filters and processes for JAX.

differentiable_assimilate(filter_: AbstractSequentialFilter, dynamics: AbstractDynamics, obs_op: AbstractObsOperator, init_ensemble: Float[Array, 'N_e N_x'], observations: Float[Array, 'T N_y'], obs_times: Float[Array, ' T'], obs_noise: lx.AbstractLinearOperator, *, inflator: AbstractInflator | None = None, t0: float | Float[Array, ''] = 0.0, adjoint: Any | None = None, checkpoint: bool = False, **analysis_extra: Any) -> AssimilationResult

Scan-based assimilation loop suitable for jax.grad / jax.jit.

Drop-in replacement for the L2 model.assimilate(...) Python loop when the caller needs a single fused XLA While (large T, short compile, optional gradient checkpointing). Observations and timestamps are stacked along a leading time axis so the loop body has a static shape.

Parameters:

Name	Type	Description	Default
filter_	AbstractSequentialFilter	Deterministic L1 filter (ETKF, EnSRF, ESTKF, EnSRF_Serial, or LETKF). Stochastic filters raise.	required
dynamics	AbstractDynamics	Forward model applied with :func:equinox.filter_vmap over the ensemble.	required
obs_op	AbstractObsOperator	Observation operator H.	required
init_ensemble	Float[Array, 'N_e N_x']	Prior ensemble (Nₑ, Nₓ).	required
observations	Float[Array, 'T N_y']	Stacked obs (T, Nᵧ).	required
obs_times	Float[Array, ' T']	Stacked timestamps (T,). The forecast for window t propagates from obs_times[t-1] (or t0 when t == 0) to obs_times[t].	required
obs_noise	AbstractLinearOperator	Observation error covariance R.	required
inflator	AbstractInflator | None	Optional deterministic inflator (MultiplicativeInflator / RTPS / RTPP). AdditiveInflator raises.	None
t0	float | Float[Array, '']	Starting time used for the first forecast window. Default 0.	0.0
adjoint	Any | None	Gradient strategy for the scan — a :class:filterax.differentiable.DirectAdjoint (default: exact, full tape), :class:filterax.differentiable.RecursiveCheckpointAdjoint (exact, recomputing), or :class:filterax.differentiable.TruncatedAdjoint (biased: gradients flow through the trailing k cycles only — O(1) backward memory in T and chaos-tolerant). The structurally identical pipekit_cycle.adjoints specs are accepted interchangeably. Forward values are identical under every strategy.	None
checkpoint	bool	Deprecated — use adjoint=RecursiveCheckpointAdjoint() instead. When True, wrap the scan body with :func:jax.checkpoint. Reduces backward-pass memory from O(T Nₑ Nₓ) to O(√T Nₑ Nₓ) at 2–3× compute cost — see §5.1 of the differentiable-DA design doc.	False
**analysis_extra	Any	Forwarded to filter_.analysis(...). Use this for filters that need extra context (LETKF takes state_coords / obs_coords here).	{}

Returns:

Type	Description
AssimilationResult	class:AssimilationResult with the same fields as
AssimilationResult	L2.assimilate(...). log_likelihoods is always populated
AssimilationResult	because every deterministic filter returns one.

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> import filterax as flx
>>> particles = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]])
>>> observations = jnp.array([[0.5], [0.4]])  # (T, N_y) = (2, 1)
>>> obs_times = jnp.array([1.0, 2.0])
>>> R = lx.DiagonalLinearOperator(0.1 * jnp.ones(1))
>>> result = flx.differentiable_assimilate(
...     flx.filters.ETKF(),
...     lambda x, t0, t1: x,  # identity dynamics
...     lambda x: x[:1],  # observe the first component
...     particles,
...     observations,
...     obs_times,
...     R,
... )
>>> result.analysis_history.shape, result.log_likelihoods.shape
((2, 3, 2), (2,))

Source code in src/filterax/_train/_differentiable.py

def differentiable_assimilate(
    filter_: AbstractSequentialFilter,
    dynamics: AbstractDynamics,
    obs_op: AbstractObsOperator,
    init_ensemble: Float[Array, "N_e N_x"],
    observations: Float[Array, "T N_y"],
    obs_times: Float[Array, " T"],
    obs_noise: lx.AbstractLinearOperator,
    *,
    inflator: AbstractInflator | None = None,
    t0: float | Float[Array, ""] = 0.0,
    adjoint: Any | None = None,
    checkpoint: bool = False,
    **analysis_extra: Any,
) -> AssimilationResult:
    r"""Scan-based assimilation loop suitable for ``jax.grad`` / ``jax.jit``.

    Drop-in replacement for the L2 ``model.assimilate(...)`` Python loop
    when the caller needs a single fused XLA ``While`` (large ``T``,
    short compile, optional gradient checkpointing). Observations and
    timestamps are stacked along a leading time axis so the loop body
    has a static shape.

    Args:
        filter_: Deterministic L1 filter (``ETKF``, ``EnSRF``, ``ESTKF``,
            ``EnSRF_Serial``, or ``LETKF``). Stochastic filters raise.
        dynamics: Forward model applied with :func:`equinox.filter_vmap`
            over the ensemble.
        obs_op: Observation operator ``H``.
        init_ensemble: Prior ensemble ``(Nₑ, Nₓ)``.
        observations: Stacked obs ``(T, Nᵧ)``.
        obs_times: Stacked timestamps ``(T,)``. The forecast for window
            ``t`` propagates from ``obs_times[t-1]`` (or ``t0`` when
            ``t == 0``) to ``obs_times[t]``.
        obs_noise: Observation error covariance ``R``.
        inflator: Optional deterministic inflator
            (``MultiplicativeInflator`` / ``RTPS`` / ``RTPP``).
            ``AdditiveInflator`` raises.
        t0: Starting time used for the first forecast window. Default 0.
        adjoint: Gradient strategy for the scan — a
            :class:`filterax.differentiable.DirectAdjoint` (default:
            exact, full tape),
            :class:`filterax.differentiable.RecursiveCheckpointAdjoint`
            (exact, recomputing), or
            :class:`filterax.differentiable.TruncatedAdjoint` (biased:
            gradients flow through the trailing ``k`` cycles only —
            O(1) backward memory in ``T`` and chaos-tolerant). The
            structurally identical ``pipekit_cycle.adjoints`` specs are
            accepted interchangeably. Forward values are identical
            under every strategy.
        checkpoint: Deprecated — use
            ``adjoint=RecursiveCheckpointAdjoint()`` instead. When
            ``True``, wrap the scan body with
            :func:`jax.checkpoint`. Reduces backward-pass memory from
            ``O(T Nₑ Nₓ)`` to ``O(√T Nₑ Nₓ)`` at 2–3× compute cost — see
            §5.1 of the differentiable-DA design doc.
        **analysis_extra: Forwarded to ``filter_.analysis(...)``. Use
            this for filters that need extra context (``LETKF`` takes
            ``state_coords`` / ``obs_coords`` here).

    Returns:
        :class:`AssimilationResult` with the same fields as
        ``L2.assimilate(...)``. ``log_likelihoods`` is always populated
        because every deterministic filter returns one.

    Examples:
        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> import filterax as flx
        >>> particles = jnp.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]])
        >>> observations = jnp.array([[0.5], [0.4]])  # (T, N_y) = (2, 1)
        >>> obs_times = jnp.array([1.0, 2.0])
        >>> R = lx.DiagonalLinearOperator(0.1 * jnp.ones(1))
        >>> result = flx.differentiable_assimilate(
        ...     flx.filters.ETKF(),
        ...     lambda x, t0, t1: x,  # identity dynamics
        ...     lambda x: x[:1],  # observe the first component
        ...     particles,
        ...     observations,
        ...     obs_times,
        ...     R,
        ... )
        >>> result.analysis_history.shape, result.log_likelihoods.shape
        ((2, 3, 2), (2,))
    """
    _reject_stochastic_components(filter_, inflator)
    if observations.shape[0] != obs_times.shape[0]:
        raise ValueError(
            "observations and obs_times must have the same leading axis; "
            f"got {observations.shape[0]} and {obs_times.shape[0]}."
        )

    # ``lax.scan`` requires the carry's input and output dtypes to match
    # exactly. ``t_prev`` enters the carry from ``t0`` and exits it from
    # the per-step ``obs_times`` element, so cast both to a common dtype
    # up front — otherwise float32 ensembles with float64 (or integer)
    # ``obs_times`` raise at trace time even though the L2 Python loop
    # would accept that combination.
    t0_arr = jnp.asarray(t0)
    time_dtype = jnp.result_type(t0_arr, obs_times)
    t0_arr = t0_arr.astype(time_dtype)
    obs_times = obs_times.astype(time_dtype)

    def step(
        carry: tuple[Float[Array, "N_e N_x"], Float[Array, ""]],
        inputs: tuple[Float[Array, " N_y"], Float[Array, ""]],
    ) -> tuple[
        tuple[Float[Array, "N_e N_x"], Float[Array, ""]],
        tuple[Float[Array, "N_e N_x"], Float[Array, "N_e N_x"], Float[Array, ""]],
    ]:
        particles, t_prev = carry
        obs_t, t_now = inputs
        forecast = _forecast(dynamics, particles, t_prev, t_now)
        result = filter_.analysis(forecast, obs_t, obs_op, obs_noise, **analysis_extra)
        analysis_particles = result.particles
        if inflator is not None:
            analysis_particles = inflator(analysis_particles, forecast)
        # ``log_likelihood`` is typed ``Array | None`` on the protocol but
        # is populated by every deterministic filter we accept. The
        # Python-level guard is resolved at trace time, so when a filter
        # ever returns ``None`` we fall back to NaN rather than feeding
        # ``None`` into ``lax.scan``.
        log_lik = result.log_likelihood
        if log_lik is None:
            log_lik = jnp.full((), jnp.nan, dtype=forecast.dtype)
        return (analysis_particles, t_now), (
            forecast,
            analysis_particles,
            log_lik,
        )

    if checkpoint:
        warnings.warn(
            "checkpoint=True is deprecated; use "
            "adjoint=RecursiveCheckpointAdjoint() instead.",
            DeprecationWarning,
            stacklevel=2,
        )
    strategy = resolve_adjoint(adjoint, checkpoint)
    strategy_name = type(strategy).__name__
    init_carry = (init_ensemble, t0_arr)

    if strategy_name == "TruncatedAdjoint":
        # Gradients flow through the trailing-k carries only: earlier
        # steps route the carry through stop_gradient (jnp.where keeps
        # the trailing window differentiable). Forward values are
        # identical to the plain scan.
        length = observations.shape[0]
        cutoff = length - strategy.k

        def truncated_step(carry, indexed):
            i, inputs = indexed
            new_carry, out = step(carry, inputs)
            cut = i < cutoff
            new_carry = jax.tree.map(
                lambda c: jnp.where(cut, jax.lax.stop_gradient(c), c), new_carry
            )
            return new_carry, out

        (final_particles, _), (forecasts, analyses, logps) = jax.lax.scan(
            truncated_step,
            init_carry,
            (jnp.arange(length), (observations, obs_times)),
        )
    else:
        step_fn = (
            jax.checkpoint(step)
            if strategy_name == "RecursiveCheckpointAdjoint"
            else step
        )
        (final_particles, _), (forecasts, analyses, logps) = jax.lax.scan(
            step_fn, init_carry, (observations, obs_times)
        )
    return AssimilationResult(
        particles=final_particles,
        forecast_history=forecasts,
        analysis_history=analyses,
        log_likelihoods=logps,
    )

Adjoint strategies

Declarative gradient policies for the scan, in the shared pipekit/diffrax vocabulary. DirectAdjoint is exact with the full tape; RecursiveCheckpointAdjoint is exact with recomputation; TruncatedAdjoint(k) cuts gradient flow more than k cycles back — biased, O(1) backward memory in T, and tolerant of chaotic gradient explosion. Structurally identical pipekit_cycle.adjoints specs are accepted interchangeably.

Differentiable data-assimilation surface.

Gradient strategies for training through the filter:

• :func:differentiable_assimilate — fixed-shape :func:jax.lax.scan whose gradient strategy is selected with adjoint=: :class:DirectAdjoint (exact, full tape), :class:RecursiveCheckpointAdjoint (exact, O(√T) memory via recomputation), or :class:TruncatedAdjoint (biased: trailing-k cycles only — O(1) memory in T and tolerant of chaotic gradient explosion). The names match the shared pipekit / diffrax vocabulary; pipekit_cycle.adjoints specs are accepted interchangeably.
• :func:road_enkf_loss_and_grad / :func:road_enkf_grad_step — ROAD-EnKF local-gradient strategy with :func:jax.lax.stop_gradient between cycles. Backward-pass memory O(Nₑ · Nₓ) independent of T; drops cross-time gradient terms.

Both refuse stochastic filters (StochasticEnKF, ETKF_Livings) and the stochastic AdditiveInflator at call time — see design_docs/features/differentiable_da.md §5.2 / §5.5.

DirectAdjoint dataclass

Exact gradients with the full forward tape stored (\(O(T)\) memory).

Source code in src/filterax/_train/_adjoints.py

@dataclass(frozen=True)
class DirectAdjoint:
    """Exact gradients with the full forward tape stored ($O(T)$ memory)."""

RecursiveCheckpointAdjoint dataclass

Exact gradients via :func:jax.checkpoint recomputation.

Attributes:

Name	Type	Description
checkpoints	int | None	Reserved for schedule control; the current implementation uses :func:jax.checkpoint's default policy and ignores this value beyond carrying it as config.

Source code in src/filterax/_train/_adjoints.py

@dataclass(frozen=True)
class RecursiveCheckpointAdjoint:
    """Exact gradients via :func:`jax.checkpoint` recomputation.

    Attributes:
        checkpoints: Reserved for schedule control; the current
            implementation uses :func:`jax.checkpoint`'s default policy
            and ignores this value beyond carrying it as config.
    """

    checkpoints: int | None = None

TruncatedAdjoint dataclass

Gradients flow through the trailing k cycles only.

Cross-cycle flow through the ensemble carry is cut more than k cycles back; per-cycle outputs keep their local gradients, so history-summed losses accumulate one local term per window (the ROAD-EnKF estimator at k=1) instead of exploding cross-window products.

Attributes:

Name	Type	Description
k	int	Number of trailing cycles that propagate gradients through the ensemble carry. Must be at least 1; k >= T is equivalent to :class:DirectAdjoint.

Source code in src/filterax/_train/_adjoints.py

@dataclass(frozen=True)
class TruncatedAdjoint:
    """Gradients flow through the trailing ``k`` cycles only.

    Cross-cycle flow through the ensemble carry is cut more than
    ``k`` cycles back; per-cycle outputs keep their local gradients,
    so history-summed losses accumulate one local term per window
    (the ROAD-EnKF estimator at ``k=1``) instead of exploding
    cross-window products.

    Attributes:
        k: Number of trailing cycles that propagate gradients through
            the ensemble carry. Must be at least 1; ``k >= T`` is
            equivalent to :class:`DirectAdjoint`.
    """

    k: int = 1

ROAD-EnKF local gradients

road_enkf_loss_and_grad evaluates the loss and its local gradient in one pass; road_enkf_grad_step wraps it with an optax optimizer update for a ready-made training step. Both live in filterax.differentiable.

Differentiable data-assimilation surface.

Gradient strategies for training through the filter:

• :func:differentiable_assimilate — fixed-shape :func:jax.lax.scan whose gradient strategy is selected with adjoint=: :class:DirectAdjoint (exact, full tape), :class:RecursiveCheckpointAdjoint (exact, O(√T) memory via recomputation), or :class:TruncatedAdjoint (biased: trailing-k cycles only — O(1) memory in T and tolerant of chaotic gradient explosion). The names match the shared pipekit / diffrax vocabulary; pipekit_cycle.adjoints specs are accepted interchangeably.
• :func:road_enkf_loss_and_grad / :func:road_enkf_grad_step — ROAD-EnKF local-gradient strategy with :func:jax.lax.stop_gradient between cycles. Backward-pass memory O(Nₑ · Nₓ) independent of T; drops cross-time gradient terms.

Both refuse stochastic filters (StochasticEnKF, ETKF_Livings) and the stochastic AdditiveInflator at call time — see design_docs/features/differentiable_da.md §5.2 / §5.5.

road_enkf_loss_and_grad(dynamics: AbstractDynamics, init_ensemble: Float[Array, 'N_e N_x'], observations: Float[Array, 'T N_y'], obs_times: Float[Array, ' T'], obs_noise: lx.AbstractLinearOperator, obs_op: AbstractObsOperator, *, filter_: AbstractSequentialFilter | None = None, inflator: AbstractInflator | None = None, t0: float | Float[Array, ''] = 0.0, **analysis_extra: Any) -> tuple[Float[Array, ''], PyTree]

Local-gradient NLL loss and dynamics gradient on a filter window.

Per-step :func:equinox.filter_value_and_grad over the forecast + analysis cycle, with :func:jax.lax.stop_gradient between cycles. Backward-pass memory is O(Nₑ · Nₓ) independent of T (cf. :func:differentiable_assimilate whose tape grows linearly in T).

Parameters:

Name	Type	Description	Default
dynamics	AbstractDynamics	Forward model. Trainable: every array leaf of this AbstractDynamics PyTree contributes a gradient.	required
init_ensemble	Float[Array, 'N_e N_x']	Prior ensemble (Nₑ, Nₓ).	required
observations	Float[Array, 'T N_y']	Stacked observation history (T, Nᵧ).	required
obs_times	Float[Array, ' T']	Stacked timestamps (T,) — the forecast for window t propagates from obs_times[t-1] (or t0 when t == 0) to obs_times[t].	required
obs_noise	AbstractLinearOperator	Observation error covariance R.	required
obs_op	AbstractObsOperator	Observation operator H.	required
filter_	AbstractSequentialFilter | None	Deterministic L1 filter (ETKF by default). Stochastic filters (StochasticEnKF, ETKF_Livings) and the stochastic AdditiveInflator are rejected at call time — their PRNG draws break smooth gradients.	None
inflator	AbstractInflator | None	Optional deterministic inflator (MultiplicativeInflator / RTPS / RTPP).	None
t0	float | Float[Array, '']	Starting time used for the first forecast window.	0.0
**analysis_extra	Any	Forwarded to filter_.analysis(...) (use for LETKF's state_coords / obs_coords).	{}

Returns:

Type	Description
Float[Array, '']	(loss, dynamics_grad) where loss is the summed
PyTree	observation-space negative log-likelihood and dynamics_grad
tuple[Float[Array, ''], PyTree]	is a PyTree of the same structure as dynamics (only its
tuple[Float[Array, ''], PyTree]	array leaves carry gradients; static / non-array fields are
tuple[Float[Array, ''], PyTree]	None-like via eqx.filter_value_and_grad).

Source code in src/filterax/_train/_road.py

def road_enkf_loss_and_grad(
    dynamics: AbstractDynamics,
    init_ensemble: Float[Array, "N_e N_x"],
    observations: Float[Array, "T N_y"],
    obs_times: Float[Array, " T"],
    obs_noise: lx.AbstractLinearOperator,
    obs_op: AbstractObsOperator,
    *,
    filter_: AbstractSequentialFilter | None = None,
    inflator: AbstractInflator | None = None,
    t0: float | Float[Array, ""] = 0.0,
    **analysis_extra: Any,
) -> tuple[Float[Array, ""], PyTree]:
    r"""Local-gradient NLL loss and dynamics gradient on a filter window.

    Per-step :func:`equinox.filter_value_and_grad` over the forecast +
    analysis cycle, with :func:`jax.lax.stop_gradient` between cycles.
    Backward-pass memory is ``O(Nₑ · Nₓ)`` independent of ``T`` (cf.
    :func:`differentiable_assimilate` whose tape grows linearly in
    ``T``).

    Args:
        dynamics: Forward model. Trainable: every array leaf of this
            ``AbstractDynamics`` PyTree contributes a gradient.
        init_ensemble: Prior ensemble ``(Nₑ, Nₓ)``.
        observations: Stacked observation history ``(T, Nᵧ)``.
        obs_times: Stacked timestamps ``(T,)`` — the forecast for window
            ``t`` propagates from ``obs_times[t-1]`` (or ``t0`` when
            ``t == 0``) to ``obs_times[t]``.
        obs_noise: Observation error covariance ``R``.
        obs_op: Observation operator ``H``.
        filter_: Deterministic L1 filter (``ETKF`` by default).
            Stochastic filters (``StochasticEnKF``, ``ETKF_Livings``)
            and the stochastic ``AdditiveInflator`` are rejected at
            call time — their PRNG draws break smooth gradients.
        inflator: Optional deterministic inflator
            (``MultiplicativeInflator`` / ``RTPS`` / ``RTPP``).
        t0: Starting time used for the first forecast window.
        **analysis_extra: Forwarded to ``filter_.analysis(...)`` (use
            for LETKF's ``state_coords`` / ``obs_coords``).

    Returns:
        ``(loss, dynamics_grad)`` where ``loss`` is the summed
        observation-space negative log-likelihood and ``dynamics_grad``
        is a PyTree of the same structure as ``dynamics`` (only its
        array leaves carry gradients; static / non-array fields are
        ``None``-like via ``eqx.filter_value_and_grad``).
    """
    # Use a deterministic default filter; import here to avoid circular import.
    if filter_ is None:
        from filterax._filters._sequential import ETKF

        filter_ = ETKF()
    _reject_stochastic_components(filter_, inflator)
    if observations.shape[0] != obs_times.shape[0]:
        raise ValueError(
            "observations and obs_times must have the same leading axis; "
            f"got {observations.shape[0]} and {obs_times.shape[0]}."
        )

    def local_step(
        dyn: AbstractDynamics,
        ensemble: Float[Array, "N_e N_x"],
        t_prev: Float[Array, ""],
        t_now: Float[Array, ""],
        obs_t: Float[Array, " N_y"],
    ) -> tuple[Float[Array, ""], Float[Array, "N_e N_x"]]:
        """One forecast + analysis with gradient tape on ``dyn`` only."""
        forecast = _forecast(dyn, ensemble, t_prev, t_now)
        result = filter_.analysis(forecast, obs_t, obs_op, obs_noise, **analysis_extra)
        analysis_particles = result.particles
        if inflator is not None:
            analysis_particles = inflator(analysis_particles, forecast)
        # Every deterministic filter returns a log-likelihood; the protocol's
        # ``Array | None`` slot exists for future stochastic variants.
        log_lik = result.log_likelihood
        if log_lik is None:
            log_lik = jnp.full((), jnp.nan, dtype=ensemble.dtype)
        return -log_lik, analysis_particles

    # ``filter_value_and_grad`` with ``has_aux`` returns the analysis
    # particles alongside the local loss, so we don't recompute the
    # forecast + analysis to advance the carry.
    grad_step = eqx.filter_value_and_grad(local_step, has_aux=True)

    # Unify the time dtype so the per-step ``t_prev`` carry stays consistent
    # regardless of caller-supplied dtypes — mirrors the same fix in
    # ``differentiable_assimilate``.
    t0_arr = jnp.asarray(t0)
    time_dtype = jnp.result_type(t0_arr, obs_times)
    t0_arr = t0_arr.astype(time_dtype)
    obs_times = obs_times.astype(time_dtype)

    # Seed the gradient accumulator with the right pytree shape — zeros
    # over *inexact* array leaves (float / complex) of ``dynamics``.
    # ``eqx.filter_value_and_grad`` only ever produces gradients for
    # inexact leaves, so we use ``eqx.is_inexact_array`` here to skip
    # integer/bool array leaves (indices, masks, PRNG-key metadata)
    # rather than seeding zero "gradients" optax would then try to
    # apply to them.
    total_grad_seed: PyTree = jax.tree.map(
        lambda leaf: jnp.zeros_like(leaf) if eqx.is_inexact_array(leaf) else None,
        eqx.filter(dynamics, eqx.is_inexact_array),
    )
    loss_seed = jnp.asarray(0.0, dtype=init_ensemble.dtype)

    # ``jax.lax.scan`` over the observation history. Each iteration runs
    # ``filter_value_and_grad`` on the per-step local loss; the resulting
    # gradient accumulates into a carry leaf of the same dynamics-shaped
    # PyTree. ``stop_gradient`` on the analysis particles between
    # iterations stops any reverse-mode tape from spanning steps — so the
    # backward-pass memory stays ``O(Nₑ · Nₓ)`` and, equally important,
    # the *forward* JIT trace stays ``O(1)`` in ``T`` (the loop body is
    # traced once and ``lax.scan`` carries it).
    def scan_step(carry, inputs):
        ensemble, total_loss, total_grad, t_prev = carry
        obs_t, t_now = inputs
        (loss_t, analysis_particles), grad_t = grad_step(
            dynamics, ensemble, t_prev, t_now, obs_t
        )
        new_total_loss = total_loss + loss_t
        new_total_grad = jax.tree.map(lambda acc, g: acc + g, total_grad, grad_t)
        new_ensemble = jax.lax.stop_gradient(analysis_particles)
        return (new_ensemble, new_total_loss, new_total_grad, t_now), None

    init_carry = (init_ensemble, loss_seed, total_grad_seed, t0_arr)
    (_, total_loss, total_grad, _), _ = jax.lax.scan(
        scan_step, init_carry, (observations, obs_times)
    )
    return total_loss, total_grad

road_enkf_grad_step(dynamics: AbstractDynamics, optimizer: optax.GradientTransformation, opt_state: optax.OptState, init_ensemble: Float[Array, 'N_e N_x'], observations: Float[Array, 'T N_y'], obs_times: Float[Array, ' T'], obs_noise: lx.AbstractLinearOperator, obs_op: AbstractObsOperator, *, filter_: AbstractSequentialFilter | None = None, inflator: AbstractInflator | None = None, t0: float | Float[Array, ''] = 0.0, **analysis_extra: Any) -> tuple[AbstractDynamics, optax.OptState, Float[Array, '']]

One ROAD-EnKF descent step composed with an :mod:optax optimizer.

Convenience wrapper that runs :func:road_enkf_loss_and_grad, threads the gradient through optimizer.update, and applies the updates to the dynamics module. Returns (new_dynamics, new_opt_state, loss) so the caller can log loss and continue the training loop.

Parameters:

Name	Type	Description	Default
dynamics	AbstractDynamics	Forward model to update.	required
optimizer	GradientTransformation	Any optax.GradientTransformation (optax.adam(1e-3), optax.chain(...), …).	required
opt_state	OptState	Optimizer state. Initialise with the inexact (gradient-eligible) leaves only — matches what the gradient update will write to:: opt_state = optimizer.init( eqx.filter(dynamics, eqx.is_inexact_array) ) Filtering on eqx.is_array instead would let optax try to write float updates onto integer / boolean metadata (indices, masks, …) and either fail or silently corrupt those fields.	required
init_ensemble	Float[Array, 'N_e N_x']	Prior ensemble (Nₑ, Nₓ).	required
observations	Float[Array, 'T N_y']	Stacked observation history (T, Nᵧ).	required
obs_times	Float[Array, ' T']	Stacked timestamps (T,).	required
obs_noise	AbstractLinearOperator	Observation error covariance R.	required
obs_op	AbstractObsOperator	Observation operator H.	required
filter_	AbstractSequentialFilter | None	Deterministic L1 filter (default ETKF).	None
inflator	AbstractInflator | None	Optional deterministic inflator.	None
t0	float | Float[Array, '']	Starting time used for the first forecast window.	0.0
**analysis_extra	Any	Forwarded to filter_.analysis(...).	{}

Returns:

Type	Description
AbstractDynamics	(new_dynamics, new_opt_state, loss) — the updated module,
OptState	the advanced optimizer state, and the loss at the current
Float[Array, '']	dynamics.

Source code in src/filterax/_train/_road.py

def road_enkf_grad_step(
    dynamics: AbstractDynamics,
    optimizer: optax.GradientTransformation,
    opt_state: optax.OptState,
    init_ensemble: Float[Array, "N_e N_x"],
    observations: Float[Array, "T N_y"],
    obs_times: Float[Array, " T"],
    obs_noise: lx.AbstractLinearOperator,
    obs_op: AbstractObsOperator,
    *,
    filter_: AbstractSequentialFilter | None = None,
    inflator: AbstractInflator | None = None,
    t0: float | Float[Array, ""] = 0.0,
    **analysis_extra: Any,
) -> tuple[AbstractDynamics, optax.OptState, Float[Array, ""]]:
    """One ROAD-EnKF descent step composed with an :mod:`optax` optimizer.

    Convenience wrapper that runs :func:`road_enkf_loss_and_grad`, threads
    the gradient through ``optimizer.update``, and applies the updates
    to the dynamics module. Returns ``(new_dynamics, new_opt_state, loss)``
    so the caller can log loss and continue the training loop.

    Args:
        dynamics: Forward model to update.
        optimizer: Any ``optax.GradientTransformation``
            (``optax.adam(1e-3)``, ``optax.chain(...)``, …).
        opt_state: Optimizer state. Initialise with the *inexact*
            (gradient-eligible) leaves only — matches what the gradient
            update will write to::

                opt_state = optimizer.init(
                    eqx.filter(dynamics, eqx.is_inexact_array)
                )

            Filtering on ``eqx.is_array`` instead would let optax try
            to write float updates onto integer / boolean metadata
            (indices, masks, …) and either fail or silently corrupt
            those fields.
        init_ensemble: Prior ensemble ``(Nₑ, Nₓ)``.
        observations: Stacked observation history ``(T, Nᵧ)``.
        obs_times: Stacked timestamps ``(T,)``.
        obs_noise: Observation error covariance ``R``.
        obs_op: Observation operator ``H``.
        filter_: Deterministic L1 filter (default ``ETKF``).
        inflator: Optional deterministic inflator.
        t0: Starting time used for the first forecast window.
        **analysis_extra: Forwarded to ``filter_.analysis(...)``.

    Returns:
        ``(new_dynamics, new_opt_state, loss)`` — the updated module,
        the advanced optimizer state, and the loss at the current
        ``dynamics``.
    """
    loss, grads = road_enkf_loss_and_grad(
        dynamics,
        init_ensemble,
        observations,
        obs_times,
        obs_noise,
        obs_op,
        filter_=filter_,
        inflator=inflator,
        t0=t0,
        **analysis_extra,
    )
    # Pass only the *inexact* (gradient-eligible) array leaves of
    # ``dynamics`` to optax — matches the ``eqx.is_inexact_array`` filter
    # used to seed the gradient accumulator in
    # ``road_enkf_loss_and_grad`` and the convention callers should use
    # when constructing ``opt_state = optimizer.init(eqx.filter(dynamics,
    # eqx.is_inexact_array))``. Filtering on ``eqx.is_array`` instead
    # would let optax try to write updates onto integer / bool leaves.
    params = eqx.filter(dynamics, eqx.is_inexact_array)
    updates, new_opt_state = optimizer.update(grads, opt_state, params)
    new_dynamics = eqx.apply_updates(dynamics, updates)
    return new_dynamics, new_opt_state, loss