Skip to content

Training & Adjoints

Training a learned solver means differentiating through an inner optimisation loop, and how you do that determines memory cost and gradient quality. The adjoint strategies here make that choice explicit and swappable — see Adjoint Methods in the Mathematical Reference for the trade-offs. Around them sit the loss functions and train/eval steps for 4DVarNet and the amortized posteriors, and the ConvLSTM gradient modulators that 4DVarNet learns in place of a hand-tuned inner optimiser.

Adjoint strategies

Implementations of vardax.adjoints (Decision D15): full backpropagation with checkpointed memory (RecursiveCheckpointAdjoint), truncated one-step gradients (OneStepAdjoint, pairing with the one_step_solve_* solver functions), and implicit differentiation at a fixed point (ImplicitAdjoint, pairing with solve_4dvarnet_1d_fixedpoint). All three are also accessible via the vardax.adjoints submodule namespace. Use assert_adjoint_calibrated to verify a cheap adjoint against the exact one before trusting it.

RecursiveCheckpointAdjoint and ImplicitAdjoint are re-exported from optimistix for one-stop import; see the optimistix documentation for their full signatures:

  • vardax.RecursiveCheckpointAdjoint — optimistix.RecursiveCheckpointAdjoint, exact reverse-mode backpropagation through the unrolled inner loop with binomial checkpointing (the default).
  • vardax.ImplicitAdjoint — optimistix.ImplicitAdjoint, implicit-function-theorem differentiation at a fixed point; pair with solve_4dvarnet_1d_fixedpoint.

KStepAdjoint(k) is vardax's own truncated adjoint — warmup under stop_gradient, then k differentiable steps; OneStepAdjoint is the k=1 alias (Bolte, Pauwels & Vaiter, NeurIPS 2023):

vardax — Modular variational data assimilation with learned components.

All public symbols are re-exported from the private _src subpackage so that user code imports from the top-level namespace:

import vardax

model = vardax.FourDVarNet1D(...)

KStepAdjoint

Bases: AbstractAdjoint

K-step differentiation: warmup under stop_gradient, then k live steps.

For a solver iteration \(z_{j+1} = T(z_j)\) run \(n\) times, the gradient flows only through the last \(k\) steps,

\[ \frac{d z_n}{d\phi} \approx \frac{\partial T^{k}(\bar z_{n-k})}{\partial\phi} \]

where \(\bar z_{n-k}\) is the iterate after the warmup, held fixed.

Algorithm
  1. Run the first \(n - k\) solver iterations under jax.lax.stop_gradient.
  2. Run the last \(k\) iterations differentiably.
  3. Backpropagate through those \(k\) iterations only, at memory cost \(O(k)\).

Use as the solver_adjoint argument to FourDVarNet1D / FourDVarNet2D:

from vardax.adjoints import KStepAdjoint

model = FourDVarNet1D(
    state_dim=N, n_time=T, ...,
    solver_adjoint=KStepAdjoint(k=3),
    key=key,
)

Raises:

Type Description
ConfigurationError

If k is below 1.

Attributes:

Name Type Description
k int

Number of trailing solver iterations that propagate gradients. Must be at least 1; values larger than n_solver_steps are clipped to a fully differentiable solve.

Examples:

>>> import vardax
>>> vardax.KStepAdjoint(k=3)
KStepAdjoint(k=3)

k=1 is the one-step estimator, and OneStepAdjoint is a subclass.

>>> isinstance(vardax.OneStepAdjoint(), vardax.KStepAdjoint)
True
>>> vardax.KStepAdjoint().k
1

A non-positive k is rejected.

>>> vardax.KStepAdjoint(k=0)
Traceback (most recent call last):
...
ConfigurationError: KStepAdjoint needs k >= 1; got k=0.
References

Bolte, J., Pauwels, E. & Vaiter, S. (2023). One-step differentiation of iterative algorithms. NeurIPS 36. arXiv:2305.13768.

Source code in src/vardax/_src/adjoints/k_step.py 
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
class KStepAdjoint(optx.AbstractAdjoint):
    r"""K-step differentiation: warmup under ``stop_gradient``, then ``k`` live steps.

    For a solver iteration $z_{j+1} = T(z_j)$ run $n$ times, the gradient
    flows only through the last $k$ steps,

    $$ \frac{d z_n}{d\phi} \approx \frac{\partial T^{k}(\bar z_{n-k})}{\partial\phi} $$

    where $\bar z_{n-k}$ is the iterate after the warmup, held fixed.

    Algorithm:
        1. Run the first $n - k$ solver iterations under
           ``jax.lax.stop_gradient``.
        2. Run the last $k$ iterations differentiably.
        3. Backpropagate through those $k$ iterations only, at memory cost
           $O(k)$.

    Use as the ``solver_adjoint`` argument to
    [`FourDVarNet1D`][vardax.FourDVarNet1D] /
    [`FourDVarNet2D`][vardax.FourDVarNet2D]:

    ```python
    from vardax.adjoints import KStepAdjoint

    model = FourDVarNet1D(
        state_dim=N, n_time=T, ...,
        solver_adjoint=KStepAdjoint(k=3),
        key=key,
    )
    ```

    Raises:
        ConfigurationError: If ``k`` is below 1.

    Attributes:
        k: Number of trailing solver iterations that propagate
            gradients. Must be at least 1; values larger than
            ``n_solver_steps`` are clipped to a fully differentiable
            solve.

    Examples:
        >>> import vardax
        >>> vardax.KStepAdjoint(k=3)
        KStepAdjoint(k=3)

        ``k=1`` is the one-step estimator, and `OneStepAdjoint` is a
        subclass.

        >>> isinstance(vardax.OneStepAdjoint(), vardax.KStepAdjoint)
        True
        >>> vardax.KStepAdjoint().k
        1

        A non-positive ``k`` is rejected.

        >>> vardax.KStepAdjoint(k=0)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        ConfigurationError: KStepAdjoint needs k >= 1; got k=0.

    References:
        Bolte, J., Pauwels, E. & Vaiter, S. (2023). One-step
        differentiation of iterative algorithms. NeurIPS 36.
        [arXiv:2305.13768](https://arxiv.org/abs/2305.13768).
    """

    k: int = 1

    def __check_init__(self) -> None:
        if self.k < 1:
            raise ConfigurationError(f"KStepAdjoint needs k >= 1; got k={self.k}.")

    def apply(
        self,
        primal_fn: Callable,
        rewrite_fn: Callable,
        inputs: PyTree,
        tags: frozenset[object],
    ) -> PyTree[Array]:
        """Not used by vardax's custom learned solver — dispatch happens
        in ``vardax._src.solver`` via ``isinstance`` on the adjoint type.

        Implementing this method for upstream ``optimistix.minimise``
        compatibility is tracked under the planned upstream
        contribution (Decision D6).

        Args:
            primal_fn: The primal solve (unused).
            rewrite_fn: The fixed-point rewrite (unused).
            inputs: The solve inputs (unused).
            tags: Optimistix tags (unused).

        Raises:
            NotImplementedError: Always; select the adjoint through
                ``FourDVarNet*(solver_adjoint=KStepAdjoint(k=...))``.
        """
        raise NotImplementedError(
            "KStepAdjoint is currently a marker / strategy selector for the "
            "FourDVarNet inner solver. Generic apply() support for "
            "optimistix.minimise is planned as part of the upstream "
            "contribution. Use it via FourDVarNet*(solver_adjoint=KStepAdjoint(k=...))."
        )

    def __repr__(self) -> str:  # pragma: no cover - cosmetic
        return f"KStepAdjoint(k={self.k})"

    def get_config(self) -> dict[str, Any]:
        return {
            "k": self.k,
        }

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> KStepAdjoint:
        return cls(**config)

apply 

apply(
    primal_fn: Callable,
    rewrite_fn: Callable,
    inputs: PyTree,
    tags: frozenset[object],
) -> PyTree[Array]

Not used by vardax's custom learned solver — dispatch happens in vardax._src.solver via isinstance on the adjoint type.

Implementing this method for upstream optimistix.minimise compatibility is tracked under the planned upstream contribution (Decision D6).

Parameters:

Name Type Description Default
primal_fn Callable

The primal solve (unused).

 required
rewrite_fn Callable

The fixed-point rewrite (unused).

 required
inputs PyTree

The solve inputs (unused).

 required
tags frozenset[object]

Optimistix tags (unused).

 required

Raises:

Type Description
NotImplementedError

Always; select the adjoint through FourDVarNet*(solver_adjoint=KStepAdjoint(k=...)).
Source code in src/vardax/_src/adjoints/k_step.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
def apply(
    self,
    primal_fn: Callable,
    rewrite_fn: Callable,
    inputs: PyTree,
    tags: frozenset[object],
) -> PyTree[Array]:
    """Not used by vardax's custom learned solver — dispatch happens
    in ``vardax._src.solver`` via ``isinstance`` on the adjoint type.

    Implementing this method for upstream ``optimistix.minimise``
    compatibility is tracked under the planned upstream
    contribution (Decision D6).

    Args:
        primal_fn: The primal solve (unused).
        rewrite_fn: The fixed-point rewrite (unused).
        inputs: The solve inputs (unused).
        tags: Optimistix tags (unused).

    Raises:
        NotImplementedError: Always; select the adjoint through
            ``FourDVarNet*(solver_adjoint=KStepAdjoint(k=...))``.
    """
    raise NotImplementedError(
        "KStepAdjoint is currently a marker / strategy selector for the "
        "FourDVarNet inner solver. Generic apply() support for "
        "optimistix.minimise is planned as part of the upstream "
        "contribution. Use it via FourDVarNet*(solver_adjoint=KStepAdjoint(k=...))."
    )

OneStepAdjoint

Bases: KStepAdjoint

One-step differentiation (Bolte et al., 2023).

Run K - 1 warmup iterations of the inner solver with jax.lax.stop_gradient, then one differentiable step. Gives O(1) memory and is exact at the fixed point of the inner iteration.

\[ \frac{d z^\star}{d\phi} \approx \partial_\phi T(\bar z_{K-1}; \phi) \]

where \(\bar z_{K-1}\) is the iterate after the warmup, held fixed.

Algorithm
  1. Run \(K - 1\) solver iterations under jax.lax.stop_gradient.
  2. Run the last iteration differentiably.
  3. Backpropagate through that single step.

Use as the solver_adjoint argument to FourDVarNet1D / FourDVarNet2D:

from vardax.adjoints import OneStepAdjoint

model = FourDVarNet1D(
    state_dim=N, n_time=T, ...,
    solver_adjoint=OneStepAdjoint(),
    key=key,
)

Examples:

>>> import vardax
>>> adjoint = vardax.OneStepAdjoint()
>>> adjoint
OneStepAdjoint()
>>> adjoint.k
1

It is KStepAdjoint with k=1.

>>> isinstance(adjoint, vardax.KStepAdjoint)
True
References

Bolte, J., Pauwels, E. & Vaiter, S. (2023). One-step differentiation of iterative algorithms. NeurIPS 36. arXiv:2305.13768.

Source code in src/vardax/_src/adjoints/one_step.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
class OneStepAdjoint(KStepAdjoint):
    r"""One-step differentiation (Bolte et al., 2023).

    Run ``K - 1`` warmup iterations of the inner solver with
    ``jax.lax.stop_gradient``, then one differentiable step. Gives
    O(1) memory and is exact at the fixed point of the inner
    iteration.

    $$ \frac{d z^\star}{d\phi} \approx \partial_\phi T(\bar z_{K-1}; \phi) $$

    where $\bar z_{K-1}$ is the iterate after the warmup, held fixed.

    Algorithm:
        1. Run $K - 1$ solver iterations under ``jax.lax.stop_gradient``.
        2. Run the last iteration differentiably.
        3. Backpropagate through that single step.

    Use as the ``solver_adjoint`` argument to
    [`FourDVarNet1D`][vardax.FourDVarNet1D] /
    [`FourDVarNet2D`][vardax.FourDVarNet2D]:

    ```python
    from vardax.adjoints import OneStepAdjoint

    model = FourDVarNet1D(
        state_dim=N, n_time=T, ...,
        solver_adjoint=OneStepAdjoint(),
        key=key,
    )
    ```

    Examples:
        >>> import vardax
        >>> adjoint = vardax.OneStepAdjoint()
        >>> adjoint
        OneStepAdjoint()
        >>> adjoint.k
        1

        It is [`KStepAdjoint`][vardax.KStepAdjoint] with ``k=1``.

        >>> isinstance(adjoint, vardax.KStepAdjoint)
        True

    References:
        Bolte, J., Pauwels, E. & Vaiter, S. (2023). One-step
        differentiation of iterative algorithms. NeurIPS 36.
        [arXiv:2305.13768](https://arxiv.org/abs/2305.13768).
    """

    k: int = 1

    def __repr__(self) -> str:  # pragma: no cover - cosmetic
        return "OneStepAdjoint()"

to_optimistix_adjoint 

to_optimistix_adjoint(spec: Any) -> AbstractAdjoint

Map an adjoint spec onto an optimistix.AbstractAdjoint.

Algorithm
  1. Return spec unchanged if it already is an optimistix.AbstractAdjoint.
  2. Otherwise match it by class name against the mapping below.
  3. Raise for names that have no inner-solve counterpart.

Mapping:

  • ImplicitAdjoint() → optx.ImplicitAdjoint() — exact at a converged fixed point, O(1) memory, one Hessian linear solve.
  • RecursiveCheckpointAdjoint(checkpoints) → optx.RecursiveCheckpointAdjoint(checkpoints) — exact, recomputing.
  • TruncatedAdjoint(k) → KStepAdjoint(k) — warmup under stop_gradient, then k differentiable steps (k=1 is OneStepAdjoint).
  • DirectAdjoint / BacksolveAdjoint → UnsupportedInputError: the first is the plain unrolled default (pass optx.RecursiveCheckpointAdjoint() or nothing), the second only exists at the dynamics layer.

Parameters:

Name Type Description Default
spec Any

A spec from pipekit_cycle.adjoints or any structurally identical object. A ready-made optx.AbstractAdjoint passes through unchanged.

 required

Returns:

Type Description
AbstractAdjoint

The corresponding optimistix adjoint instance.

Raises:

Type Description
UnsupportedInputError

For layer-inappropriate or unrecognised specs.

Examples:

>>> import optimistix as optx
>>> from pipekit_cycle.adjoints import TruncatedAdjoint
>>> from vardax.adjoints import to_optimistix_adjoint
>>> to_optimistix_adjoint(TruncatedAdjoint(k=3))
KStepAdjoint(k=3)
>>> isinstance(
...     to_optimistix_adjoint(optx.ImplicitAdjoint()), optx.ImplicitAdjoint
... )
True

A pipekit TruncatedAdjoint(k=1) is the one-step adjoint.

>>> import vardax
>>> adjoint = to_optimistix_adjoint(TruncatedAdjoint(k=1))
>>> isinstance(adjoint, vardax.OneStepAdjoint)
False
>>> isinstance(adjoint, vardax.KStepAdjoint)
True

A spec with no inner-solve meaning is rejected.

>>> from pipekit_cycle.adjoints import DirectAdjoint
>>> to_optimistix_adjoint(DirectAdjoint())
Traceback (most recent call last):
...
UnsupportedInputError: DirectAdjoint does not apply at the inner-solve layer
Source code in src/vardax/_src/adjoints/mapping.py 
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
def to_optimistix_adjoint(spec: Any) -> optx.AbstractAdjoint:
    """Map an adjoint spec onto an ``optimistix.AbstractAdjoint``.

    Algorithm:
        1. Return ``spec`` unchanged if it already is an
           ``optimistix.AbstractAdjoint``.
        2. Otherwise match it by class name against the mapping below.
        3. Raise for names that have no inner-solve counterpart.

    Mapping:

    - ``ImplicitAdjoint()`` → ``optx.ImplicitAdjoint()`` — exact at a
      converged fixed point, O(1) memory, one Hessian linear solve.
    - ``RecursiveCheckpointAdjoint(checkpoints)`` →
      ``optx.RecursiveCheckpointAdjoint(checkpoints)`` — exact,
      recomputing.
    - ``TruncatedAdjoint(k)`` →
      [`KStepAdjoint(k)`][vardax.adjoints.KStepAdjoint] — warmup under
      ``stop_gradient``, then ``k`` differentiable steps (``k=1`` is
      [`OneStepAdjoint`][vardax.adjoints.OneStepAdjoint]).
    - ``DirectAdjoint`` / ``BacksolveAdjoint`` →
      [`UnsupportedInputError`][vardax.UnsupportedInputError]: the
      first is the plain unrolled default (pass
      ``optx.RecursiveCheckpointAdjoint()`` or nothing), the second
      only exists at the dynamics layer.

    Args:
        spec: A spec from ``pipekit_cycle.adjoints`` or any
            structurally identical object. A ready-made
            ``optx.AbstractAdjoint`` passes through unchanged.

    Returns:
        The corresponding optimistix adjoint instance.

    Raises:
        UnsupportedInputError: For layer-inappropriate or unrecognised specs.

    Examples:
        >>> import optimistix as optx
        >>> from pipekit_cycle.adjoints import TruncatedAdjoint
        >>> from vardax.adjoints import to_optimistix_adjoint
        >>> to_optimistix_adjoint(TruncatedAdjoint(k=3))
        KStepAdjoint(k=3)
        >>> isinstance(
        ...     to_optimistix_adjoint(optx.ImplicitAdjoint()), optx.ImplicitAdjoint
        ... )
        True

        A pipekit ``TruncatedAdjoint(k=1)`` is the one-step adjoint.

        >>> import vardax
        >>> adjoint = to_optimistix_adjoint(TruncatedAdjoint(k=1))
        >>> isinstance(adjoint, vardax.OneStepAdjoint)
        False
        >>> isinstance(adjoint, vardax.KStepAdjoint)
        True

        A spec with no inner-solve meaning is rejected.

        >>> from pipekit_cycle.adjoints import DirectAdjoint
        >>> to_optimistix_adjoint(DirectAdjoint())  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        UnsupportedInputError: DirectAdjoint does not apply at the inner-solve layer
    """
    if isinstance(spec, optx.AbstractAdjoint):
        return spec
    name = type(spec).__name__
    if name == "ImplicitAdjoint":
        return optx.ImplicitAdjoint()
    if name == "RecursiveCheckpointAdjoint":
        return optx.RecursiveCheckpointAdjoint(
            checkpoints=getattr(spec, "checkpoints", None)
        )
    if name == "TruncatedAdjoint":
        return KStepAdjoint(k=getattr(spec, "k", 1))
    if name in ("DirectAdjoint", "BacksolveAdjoint"):
        raise UnsupportedInputError(
            f"{name} does not apply at the inner-solve layer: DirectAdjoint is "
            "the plain unrolled default (use optx.RecursiveCheckpointAdjoint() "
            "or omit solver_adjoint), and BacksolveAdjoint only exists at the "
            "dynamics layer (see pipekit_jax.DiffraxForwardModel)."
        )
    raise UnsupportedInputError(f"Unrecognised adjoint spec: {spec!r}")

Gradient modulators

The learned components of 4DVarNet: ConvLSTM cells that map the raw variational-cost gradient to a descent update, satisfying the GradModulator Protocol. Their recurrent state is carried in the LSTMState1D / LSTMState2D containers.

vardax — Modular variational data assimilation with learned components.

All public symbols are re-exported from the private _src subpackage so that user code imports from the top-level namespace:

import vardax

model = vardax.FourDVarNet1D(...)

ConvLSTMGradMod1D

Bases: Module

1-D ConvLSTM-based gradient modulator.

Accepts the concatenation of the current state and its gradient as input and produces a modulated gradient update (and updated LSTM state). With \(g = \nabla J(u)\) and previous memory \((h, c)\) the gates are

\[ [i; f; \tilde c; o] = W_x [g; u] + W_h h, \quad c' = \sigma(f) \odot c + \sigma(i) \odot \tanh(\tilde c), \quad h' = \sigma(o) \odot \tanh(c'), \]

and the update is \(W_o h'\).

Algorithm
  1. Stack the gradient over the state along the channel axis and apply the input convolution; add the convolution of the hidden state.
  2. Split into the input, forget, candidate and output gates.
  3. Update the cell state \(c'\) and the hidden state \(h'\).
  4. Map \(h'\) back to the state channels with the output convolution.

Parameters:

Name Type Description Default
state_channels int

Number of channels in the state / gradient (T).

 required
hidden_dim int

Number of hidden channels in the LSTM.

 required
kernel_size int

1-D convolution kernel size.

 3
key PRNGKeyArray

PRNG key for the convolution weights.

 required

Examples:

The update has the shape of the gradient, the memory the shape of the hidden state:

>>> import jax
>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> mod = vdx.ConvLSTMGradMod1D(
...     state_channels=3, hidden_dim=5, key=jax.random.PRNGKey(0)
... )
>>> grad = jnp.ones((2, 3, 4))
>>> update, lstm = mod(grad, grad, vdx.LSTMState1D.zeros(2, 5, 4))
>>> update.shape, lstm.hidden.shape
((2, 3, 4), (2, 5, 4))

The memory is carried across solver steps, so a second call with the returned state gives a different update:

>>> update2, _ = mod(grad, grad, lstm)
>>> bool(jnp.allclose(update, update2))
False
Source code in src/vardax/_src/grad_mod.py 
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
class ConvLSTMGradMod1D(eqx.Module):
    r"""1-D ConvLSTM-based gradient modulator.

    Accepts the concatenation of the current state and its gradient as input
    and produces a modulated gradient update (and updated LSTM state).  With
    $g = \nabla J(u)$ and previous memory $(h, c)$ the gates are

    $$
    [i; f; \tilde c; o] = W_x [g; u] + W_h h, \quad
    c' = \sigma(f) \odot c + \sigma(i) \odot \tanh(\tilde c), \quad
    h' = \sigma(o) \odot \tanh(c'),
    $$

    and the update is $W_o h'$.

    Algorithm:
        1. Stack the gradient over the state along the channel axis and apply
           the input convolution; add the convolution of the hidden state.
        2. Split into the input, forget, candidate and output gates.
        3. Update the cell state $c'$ and the hidden state $h'$.
        4. Map $h'$ back to the state channels with the output convolution.

    Args:
        state_channels: Number of channels in the state / gradient (``T``).
        hidden_dim: Number of hidden channels in the LSTM.
        kernel_size: 1-D convolution kernel size.
        key: PRNG key for the convolution weights.

    Examples:
        The update has the shape of the gradient, the memory the shape of the
        hidden state:

        >>> import jax
        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> mod = vdx.ConvLSTMGradMod1D(
        ...     state_channels=3, hidden_dim=5, key=jax.random.PRNGKey(0)
        ... )
        >>> grad = jnp.ones((2, 3, 4))
        >>> update, lstm = mod(grad, grad, vdx.LSTMState1D.zeros(2, 5, 4))
        >>> update.shape, lstm.hidden.shape
        ((2, 3, 4), (2, 5, 4))

        The memory is carried across solver steps, so a second call with the
        returned state gives a different update:

        >>> update2, _ = mod(grad, grad, lstm)
        >>> bool(jnp.allclose(update, update2))
        False
    """

    state_channels: int = eqx.field(static=True)
    hidden_dim: int = eqx.field(static=True)
    kernel_size: int = eqx.field(static=True)
    _gates_input_conv: eqx.nn.Conv1d
    _gates_hidden_conv: eqx.nn.Conv1d
    _output_conv: eqx.nn.Conv1d

    def __init__(
        self,
        state_channels: int,
        hidden_dim: int,
        kernel_size: int = 3,
        *,
        key: PRNGKeyArray,
    ) -> None:
        self.state_channels = state_channels
        self.hidden_dim = hidden_dim
        self.kernel_size = kernel_size
        key_input, key_hidden, key_output = jax.random.split(key, 3)
        pad = kernel_size // 2
        self._gates_input_conv = eqx.nn.Conv1d(
            in_channels=2 * state_channels,
            out_channels=4 * hidden_dim,
            kernel_size=kernel_size,
            padding=pad,
            key=key_input,
        )
        self._gates_hidden_conv = eqx.nn.Conv1d(
            in_channels=hidden_dim,
            out_channels=4 * hidden_dim,
            kernel_size=kernel_size,
            padding=pad,
            key=key_hidden,
        )
        self._output_conv = eqx.nn.Conv1d(
            in_channels=hidden_dim,
            out_channels=state_channels,
            kernel_size=kernel_size,
            padding=pad,
            key=key_output,
        )

    def __call__(
        self,
        grad: Float[Array, "B T N"],
        state: Float[Array, "B T N"],
        lstm_state: LSTMState1D,
    ) -> tuple[Float[Array, "B T N"], LSTMState1D]:
        """Forward pass.

        Args:
            grad: Gradient of the variational cost with respect to the state.
            state: Current state estimate.
            lstm_state: Current LSTM hidden/cell state with ``hidden_dim`` channels.

        Returns:
            The modulated gradient update and the new LSTM state.
        """
        # eqx.nn.Conv1d operates on (channels, spatial). Our arrays are already
        # (B, channels, spatial) with T=channels for grad/state and H=channels
        # for the LSTM state — vmap over the leading batch dim.

        def _forward(
            grad_sample: Float[Array, "T N"],
            state_sample: Float[Array, "T N"],
            hidden_sample: Float[Array, "H N"],
            cell_sample: Float[Array, "H N"],
        ) -> tuple[Float[Array, "T N"], Float[Array, "H N"], Float[Array, "H N"]]:
            # the input channels are the gradient stacked over the state
            stacked_input = einx.id("a n, b n -> (a + b) n", grad_sample, state_sample)
            # gates = W_x [∇J; u] + W_h h
            gates = self._gates_input_conv(stacked_input) + self._gates_hidden_conv(
                hidden_sample
            )
            # [i; f; g; o] = split(gates)
            input_gate, forget_gate, cell_candidate, output_gate = einx.id(
                "(gate h) n -> gate h n", gates, gate=4
            )
            input_gate = jax.nn.sigmoid(input_gate)
            forget_gate = jax.nn.sigmoid(forget_gate)
            cell_candidate = jnp.tanh(cell_candidate)
            output_gate = jax.nn.sigmoid(output_gate)
            # c' = f ⊙ c + i ⊙ g
            new_cell = forget_gate * cell_sample + input_gate * cell_candidate
            # h' = o ⊙ tanh(c')
            new_hidden = output_gate * jnp.tanh(new_cell)
            update = self._output_conv(new_hidden)
            return update, new_hidden, new_cell

        update, new_hidden, new_cell = jax.vmap(_forward)(
            grad, state, lstm_state.hidden, lstm_state.cell
        )
        return update, LSTMState1D(hidden=new_hidden, cell=new_cell)

ConvLSTMGradMod2D

Bases: Module

2-D ConvLSTM-based gradient modulator.

Same recurrence as ConvLSTMGradMod1D, with 2-D convolutions: with \(g = \nabla J(u)\) and memory \((h, c)\),

\[ [i; f; \tilde c; o] = W_x [g; u] + W_h h, \quad c' = \sigma(f) \odot c + \sigma(i) \odot \tanh(\tilde c), \quad h' = \sigma(o) \odot \tanh(c'), \]

and the update is \(W_o h'\).

Algorithm
  1. Stack the gradient over the state along the channel axis and apply the input convolution; add the convolution of the hidden state.
  2. Split into the input, forget, candidate and output gates.
  3. Update the cell state \(c'\) and the hidden state \(h'\).
  4. Map \(h'\) back to the state channels with the output convolution.

Parameters:

Name Type Description Default
state_channels int

Number of time channels in the state / gradient.

 required
hidden_dim int

Number of hidden channels in the LSTM.

 required
kernel_size int

2-D convolution kernel size.

 3
key PRNGKeyArray

PRNG key for the convolution weights.

 required

Examples:

The update has the shape of the gradient:

>>> import jax
>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> mod = vdx.ConvLSTMGradMod2D(
...     state_channels=3, hidden_dim=2, key=jax.random.PRNGKey(0)
... )
>>> grad = jnp.ones((1, 3, 4, 5))
>>> update, lstm = mod(grad, grad, vdx.LSTMState2D.zeros(1, 2, 4, 5))
>>> update.shape, lstm.cell.shape
((1, 3, 4, 5), (1, 2, 4, 5))

A zero gradient and zero state still yield a finite update:

>>> zeros = jnp.zeros((1, 3, 4, 5))
>>> out, _ = mod(zeros, zeros, vdx.LSTMState2D.zeros(1, 2, 4, 5))
>>> bool(jnp.all(jnp.isfinite(out)))
True
Source code in src/vardax/_src/grad_mod.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
class ConvLSTMGradMod2D(eqx.Module):
    r"""2-D ConvLSTM-based gradient modulator.

    Same recurrence as [`ConvLSTMGradMod1D`][vardax.ConvLSTMGradMod1D], with 2-D
    convolutions: with $g = \nabla J(u)$ and memory $(h, c)$,

    $$
    [i; f; \tilde c; o] = W_x [g; u] + W_h h, \quad
    c' = \sigma(f) \odot c + \sigma(i) \odot \tanh(\tilde c), \quad
    h' = \sigma(o) \odot \tanh(c'),
    $$

    and the update is $W_o h'$.

    Algorithm:
        1. Stack the gradient over the state along the channel axis and apply
           the input convolution; add the convolution of the hidden state.
        2. Split into the input, forget, candidate and output gates.
        3. Update the cell state $c'$ and the hidden state $h'$.
        4. Map $h'$ back to the state channels with the output convolution.

    Args:
        state_channels: Number of time channels in the state / gradient.
        hidden_dim: Number of hidden channels in the LSTM.
        kernel_size: 2-D convolution kernel size.
        key: PRNG key for the convolution weights.

    Examples:
        The update has the shape of the gradient:

        >>> import jax
        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> mod = vdx.ConvLSTMGradMod2D(
        ...     state_channels=3, hidden_dim=2, key=jax.random.PRNGKey(0)
        ... )
        >>> grad = jnp.ones((1, 3, 4, 5))
        >>> update, lstm = mod(grad, grad, vdx.LSTMState2D.zeros(1, 2, 4, 5))
        >>> update.shape, lstm.cell.shape
        ((1, 3, 4, 5), (1, 2, 4, 5))

        A zero gradient and zero state still yield a finite update:

        >>> zeros = jnp.zeros((1, 3, 4, 5))
        >>> out, _ = mod(zeros, zeros, vdx.LSTMState2D.zeros(1, 2, 4, 5))
        >>> bool(jnp.all(jnp.isfinite(out)))
        True
    """

    state_channels: int = eqx.field(static=True)
    hidden_dim: int = eqx.field(static=True)
    kernel_size: int = eqx.field(static=True)
    _gates_input_conv: eqx.nn.Conv2d
    _gates_hidden_conv: eqx.nn.Conv2d
    _output_conv: eqx.nn.Conv2d

    def __init__(
        self,
        state_channels: int,
        hidden_dim: int,
        kernel_size: int = 3,
        *,
        key: PRNGKeyArray,
    ) -> None:
        self.state_channels = state_channels
        self.hidden_dim = hidden_dim
        self.kernel_size = kernel_size
        key_input, key_hidden, key_output = jax.random.split(key, 3)
        pad = kernel_size // 2
        self._gates_input_conv = eqx.nn.Conv2d(
            in_channels=2 * state_channels,
            out_channels=4 * hidden_dim,
            kernel_size=kernel_size,
            padding=pad,
            key=key_input,
        )
        self._gates_hidden_conv = eqx.nn.Conv2d(
            in_channels=hidden_dim,
            out_channels=4 * hidden_dim,
            kernel_size=kernel_size,
            padding=pad,
            key=key_hidden,
        )
        self._output_conv = eqx.nn.Conv2d(
            in_channels=hidden_dim,
            out_channels=state_channels,
            kernel_size=kernel_size,
            padding=pad,
            key=key_output,
        )

    def __call__(
        self,
        grad: Float[Array, "B T H W"],
        state: Float[Array, "B T H W"],
        lstm_state: LSTMState2D,
    ) -> tuple[Float[Array, "B T H W"], LSTMState2D]:
        """Forward pass.

        Args:
            grad: Gradient of the variational cost with respect to the state.
            state: Current state estimate.
            lstm_state: Current LSTM hidden/cell state.

        Returns:
            The modulated gradient update and the new LSTM state.
        """

        def _forward(
            grad_sample: Float[Array, "T H W"],
            state_sample: Float[Array, "T H W"],
            hidden_sample: Float[Array, "H_dim H W"],
            cell_sample: Float[Array, "H_dim H W"],
        ) -> tuple[
            Float[Array, "T H W"],
            Float[Array, "H_dim H W"],
            Float[Array, "H_dim H W"],
        ]:
            # the input channels are the gradient stacked over the state
            stacked_input = einx.id(
                "a h w, b h w -> (a + b) h w", grad_sample, state_sample
            )
            # gates = W_x [∇J; u] + W_h h
            gates = self._gates_input_conv(stacked_input) + self._gates_hidden_conv(
                hidden_sample
            )
            # [i; f; g; o] = split(gates)
            input_gate, forget_gate, cell_candidate, output_gate = einx.id(
                "(gate c) h w -> gate c h w", gates, gate=4
            )
            input_gate = jax.nn.sigmoid(input_gate)
            forget_gate = jax.nn.sigmoid(forget_gate)
            cell_candidate = jnp.tanh(cell_candidate)
            output_gate = jax.nn.sigmoid(output_gate)
            # c' = f ⊙ c + i ⊙ g
            new_cell = forget_gate * cell_sample + input_gate * cell_candidate
            # h' = o ⊙ tanh(c')
            new_hidden = output_gate * jnp.tanh(new_cell)
            update = self._output_conv(new_hidden)
            return update, new_hidden, new_cell

        update, new_hidden, new_cell = jax.vmap(_forward)(
            grad, state, lstm_state.hidden, lstm_state.cell
        )
        return update, LSTMState2D(hidden=new_hidden, cell=new_cell)

Losses & steps

Outer-loop training: reconstruction-based losses for 4DVarNet, the negative-log-likelihood loss for amortized posteriors, and the optax-driven train_step / amortized_train_step / eval_step that consume them.

vardax — Modular variational data assimilation with learned components.

All public symbols are re-exported from the private _src subpackage so that user code imports from the top-level namespace:

import vardax

model = vardax.FourDVarNet1D(...)

reconstruction_loss 

reconstruction_loss(
    pred: Float[Array, ...], target: Float[Array, ...]
) -> Float[Array, ""]

Mean-squared reconstruction loss.

\[ \mathcal{L}(\hat x, x) = \frac{1}{n} \sum_i (\hat x_i - x_i)^2 \]

Parameters:

Name Type Description Default
pred Float[Array, ...]

Model predictions \(\hat x\).

 required
target Float[Array, ...]

Ground-truth targets \(x\), broadcastable against pred.

 required

Returns:

Type Description
Float[Array, '']

Scalar mean-squared error.

Examples:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> float(vdx.reconstruction_loss(jnp.array([1.0, 3.0]), jnp.array([1.0, 1.0])))
2.0

A perfect prediction has zero loss, whatever the shape:

>>> x = jnp.ones((2, 3, 4))
>>> float(vdx.reconstruction_loss(x, x))
0.0
Source code in src/vardax/_src/training.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def reconstruction_loss(
    pred: Float[Array, ...],
    target: Float[Array, ...],
) -> Float[Array, ""]:
    r"""Mean-squared reconstruction loss.

    $$
    \mathcal{L}(\hat x, x) = \frac{1}{n} \sum_i (\hat x_i - x_i)^2
    $$

    Args:
        pred: Model predictions $\hat x$.
        target: Ground-truth targets $x$, broadcastable against ``pred``.

    Returns:
        Scalar mean-squared error.

    Examples:
        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> float(vdx.reconstruction_loss(jnp.array([1.0, 3.0]), jnp.array([1.0, 1.0])))
        2.0

        A perfect prediction has zero loss, whatever the shape:

        >>> x = jnp.ones((2, 3, 4))
        >>> float(vdx.reconstruction_loss(x, x))
        0.0
    """
    return jnp.mean((pred - target) ** 2)

train_loss_fn 

train_loss_fn(
    model: Any, batch: Batch1D | Batch2D
) -> Float[Array, ""]

Compute the training loss for a single batch.

\[ \mathcal{L} = \operatorname{mean}\big((\mathcal{M}(b) - x_\text{target})^2\big) \]
Algorithm
  1. Run the model on the batch to obtain the prediction.
  2. Compare it with batch.target using reconstruction_loss.

Parameters:

Name Type Description Default
model Any

Equinox module implementing __call__(batch) -> prediction.

 required
batch Batch1D | Batch2D

Training batch (must have a non-None target).

 required

Returns:

Type Description
Float[Array, '']

Scalar reconstruction loss.

Raises:

Type Description
MissingInputError

If batch.target is None.

Examples:

>>> import jax
>>> import jax.numpy as jnp
>>> import optax
>>> import vardax as vdx
>>> model = vdx.FourDVarNet1D(
...     state_dim=4,
...     n_time=3,
...     latent_dim=2,
...     hidden_dim=3,
...     n_solver_steps=1,
...     key=jax.random.PRNGKey(0),
... )
>>> ones = jnp.ones((2, 3, 4))
>>> batch = vdx.Batch1D(input=ones, mask=ones, target=jnp.zeros((2, 3, 4)))
>>> vdx.train_loss_fn(model, batch).shape
()

It differentiates with respect to the model, as train_step does:

>>> import equinox as eqx
>>> loss, grads = eqx.filter_value_and_grad(vdx.train_loss_fn)(model, batch)
>>> bool(jnp.isfinite(loss))
True

A batch without a target is rejected:

>>> unlabelled = vdx.Batch1D(input=ones, mask=ones)
>>> vdx.train_loss_fn(model, unlabelled)
Traceback (most recent call last):
...
MissingInputError: train_loss_fn requires batch.target to be set.
Source code in src/vardax/_src/training.py 
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def train_loss_fn(
    model: Any,
    batch: Batch1D | Batch2D,
) -> Float[Array, ""]:
    r"""Compute the training loss for a single batch.

    $$
    \mathcal{L} = \operatorname{mean}\big((\mathcal{M}(b) - x_\text{target})^2\big)
    $$

    Algorithm:
        1. Run the model on the batch to obtain the prediction.
        2. Compare it with ``batch.target`` using
           [`reconstruction_loss`][vardax.reconstruction_loss].

    Args:
        model: Equinox module implementing ``__call__(batch) -> prediction``.
        batch: Training batch (must have a non-``None`` ``target``).

    Returns:
        Scalar reconstruction loss.

    Raises:
        MissingInputError: If ``batch.target`` is ``None``.

    Examples:
        >>> import jax
        >>> import jax.numpy as jnp
        >>> import optax
        >>> import vardax as vdx
        >>> model = vdx.FourDVarNet1D(
        ...     state_dim=4,
        ...     n_time=3,
        ...     latent_dim=2,
        ...     hidden_dim=3,
        ...     n_solver_steps=1,
        ...     key=jax.random.PRNGKey(0),
        ... )
        >>> ones = jnp.ones((2, 3, 4))
        >>> batch = vdx.Batch1D(input=ones, mask=ones, target=jnp.zeros((2, 3, 4)))
        >>> vdx.train_loss_fn(model, batch).shape
        ()

        It differentiates with respect to the model, as ``train_step`` does:

        >>> import equinox as eqx
        >>> loss, grads = eqx.filter_value_and_grad(vdx.train_loss_fn)(model, batch)
        >>> bool(jnp.isfinite(loss))
        True

        A batch without a target is rejected:

        >>> unlabelled = vdx.Batch1D(input=ones, mask=ones)
        >>> vdx.train_loss_fn(model, unlabelled)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        MissingInputError: train_loss_fn requires batch.target to be set.
    """
    pred = model(batch)
    target = batch.target
    if target is None:
        raise MissingInputError("train_loss_fn requires batch.target to be set.")
    return reconstruction_loss(pred, target)

train_step 

train_step(
    model: Any,
    batch: Batch1D | Batch2D,
    optimizer: GradientTransformation,
    opt_state: OptState,
) -> tuple[Any, OptState, Float[Array, ""]]

Perform a single training step (forward + backward + update).

This is the correctness-critical primitive: gradients flow through the FourDVarNet inner solver according to whichever differentiation strategy (the model's solver_adjoint) the model is configured with. Users should compose this primitive into their training loop (notebook-level or pipekit_train.TrainingLoop) rather than reimplementing it.

Algorithm
  1. Evaluate train_loss_fn and its gradient with respect to the model parameters.
  2. Turn the gradient into an update with the optax optimiser.
  3. Apply the update to the model.

Parameters:

Name Type Description Default
model Any

Equinox module to optimise.

 required
batch Batch1D | Batch2D

Training batch.

 required
optimizer GradientTransformation

Optax gradient transformation (e.g. optax.adam(1e-3)).

 required
opt_state OptState

Current optimiser state.

 required

Returns:

Type Description
Any

The updated model, the updated optimiser state and the scalar loss
OptState

(evaluated before the update).

Raises:

Type Description
MissingInputError

If batch.target is None.

Examples:

>>> import jax
>>> import jax.numpy as jnp
>>> import equinox as eqx
>>> import optax
>>> import vardax as vdx
>>> model = vdx.FourDVarNet1D(
...     state_dim=4,
...     n_time=3,
...     latent_dim=2,
...     hidden_dim=3,
...     n_solver_steps=1,
...     key=jax.random.PRNGKey(0),
... )
>>> ones = jnp.ones((2, 3, 4))
>>> batch = vdx.Batch1D(input=ones, mask=ones, target=jnp.zeros((2, 3, 4)))
>>> optimizer = optax.adam(1e-2)
>>> opt_state = optimizer.init(eqx.filter(model, eqx.is_array))
>>> new_model, opt_state, loss = vdx.train_step(
...     model, batch, optimizer, opt_state
... )
>>> loss.shape
()

The step returns the pre-update loss, so it equals eval_step on the old model, and the parameters have moved:

>>> bool(jnp.allclose(loss, vdx.eval_step(model, batch)))
True
>>> bool(eqx.tree_equal(new_model, model))
False
Source code in src/vardax/_src/training.py 
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
@eqx.filter_jit
def train_step(
    model: Any,
    batch: Batch1D | Batch2D,
    optimizer: optax.GradientTransformation,
    opt_state: optax.OptState,
) -> tuple[Any, optax.OptState, Float[Array, ""]]:
    """Perform a single training step (forward + backward + update).

    This is the correctness-critical primitive: gradients flow through
    the FourDVarNet inner solver according to whichever differentiation
    strategy (the model's ``solver_adjoint``) the model is configured with.
    Users should compose this primitive into their training loop
    (notebook-level or ``pipekit_train.TrainingLoop``) rather than
    reimplementing it.

    Algorithm:
        1. Evaluate [`train_loss_fn`][vardax.train_loss_fn] and its gradient with
           respect to the model parameters.
        2. Turn the gradient into an update with the optax optimiser.
        3. Apply the update to the model.

    Args:
        model: Equinox module to optimise.
        batch: Training batch.
        optimizer: Optax gradient transformation (e.g. ``optax.adam(1e-3)``).
        opt_state: Current optimiser state.

    Returns:
        The updated model, the updated optimiser state and the scalar loss
        (evaluated before the update).

    Raises:
        MissingInputError: If ``batch.target`` is ``None``.

    Examples:
        >>> import jax
        >>> import jax.numpy as jnp
        >>> import equinox as eqx
        >>> import optax
        >>> import vardax as vdx
        >>> model = vdx.FourDVarNet1D(
        ...     state_dim=4,
        ...     n_time=3,
        ...     latent_dim=2,
        ...     hidden_dim=3,
        ...     n_solver_steps=1,
        ...     key=jax.random.PRNGKey(0),
        ... )
        >>> ones = jnp.ones((2, 3, 4))
        >>> batch = vdx.Batch1D(input=ones, mask=ones, target=jnp.zeros((2, 3, 4)))
        >>> optimizer = optax.adam(1e-2)
        >>> opt_state = optimizer.init(eqx.filter(model, eqx.is_array))
        >>> new_model, opt_state, loss = vdx.train_step(
        ...     model, batch, optimizer, opt_state
        ... )
        >>> loss.shape
        ()

        The step returns the pre-update loss, so it equals ``eval_step`` on the
        old model, and the parameters have moved:

        >>> bool(jnp.allclose(loss, vdx.eval_step(model, batch)))
        True
        >>> bool(eqx.tree_equal(new_model, model))
        False
    """
    loss, grads = eqx.filter_value_and_grad(train_loss_fn)(model, batch)
    updates, opt_state = optimizer.update(grads, opt_state, model)
    model = eqx.apply_updates(model, updates)
    return model, opt_state, loss

eval_step 

eval_step(
    model: Any, batch: Batch1D | Batch2D
) -> Float[Array, ""]

Compute the evaluation loss for a single batch (no gradient).

Parameters:

Name Type Description Default
model Any

Equinox module implementing __call__(batch) -> prediction.

 required
batch Batch1D | Batch2D

Evaluation batch (must have a non-None target).

 required

Returns:

Type Description
Float[Array, '']

Scalar reconstruction loss.

Raises:

Type Description
MissingInputError

If batch.target is None.

Examples:

>>> import jax
>>> import jax.numpy as jnp
>>> import optax
>>> import vardax as vdx
>>> model = vdx.FourDVarNet1D(
...     state_dim=4,
...     n_time=3,
...     latent_dim=2,
...     hidden_dim=3,
...     n_solver_steps=1,
...     key=jax.random.PRNGKey(0),
... )
>>> ones = jnp.ones((2, 3, 4))
>>> batch = vdx.Batch1D(input=ones, mask=ones, target=jnp.zeros((2, 3, 4)))
>>> vdx.eval_step(model, batch).shape
()

It agrees with the training loss on the same batch:

>>> bool(
...     jnp.allclose(
...         vdx.eval_step(model, batch), vdx.train_loss_fn(model, batch)
...     )
... )
True

A batch without a target is rejected:

>>> vdx.eval_step(
...     model, vdx.Batch1D(input=ones, mask=ones)
... )
Traceback (most recent call last):
...
MissingInputError: eval_step requires batch.target to be set.
Source code in src/vardax/_src/training.py 
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
@eqx.filter_jit
def eval_step(
    model: Any,
    batch: Batch1D | Batch2D,
) -> Float[Array, ""]:
    """Compute the evaluation loss for a single batch (no gradient).

    Args:
        model: Equinox module implementing ``__call__(batch) -> prediction``.
        batch: Evaluation batch (must have a non-``None`` ``target``).

    Returns:
        Scalar reconstruction loss.

    Raises:
        MissingInputError: If ``batch.target`` is ``None``.

    Examples:
        >>> import jax
        >>> import jax.numpy as jnp
        >>> import optax
        >>> import vardax as vdx
        >>> model = vdx.FourDVarNet1D(
        ...     state_dim=4,
        ...     n_time=3,
        ...     latent_dim=2,
        ...     hidden_dim=3,
        ...     n_solver_steps=1,
        ...     key=jax.random.PRNGKey(0),
        ... )
        >>> ones = jnp.ones((2, 3, 4))
        >>> batch = vdx.Batch1D(input=ones, mask=ones, target=jnp.zeros((2, 3, 4)))
        >>> vdx.eval_step(model, batch).shape
        ()

        It agrees with the training loss on the same batch:

        >>> bool(
        ...     jnp.allclose(
        ...         vdx.eval_step(model, batch), vdx.train_loss_fn(model, batch)
        ...     )
        ... )
        True

        A batch without a target is rejected:

        >>> vdx.eval_step(
        ...     model, vdx.Batch1D(input=ones, mask=ones)
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        MissingInputError: eval_step requires batch.target to be set.
    """
    pred = model(batch)
    target = batch.target
    if target is None:
        raise MissingInputError("eval_step requires batch.target to be set.")
    return reconstruction_loss(pred, target)

amortized_nll_loss_fn 

amortized_nll_loss_fn(
    model: Any, batch: Batch1D | Batch2D
) -> Float[Array, ""]

Negative log-likelihood for amortized inference (Epic 8).

For AmortizedPosterior with flow / regression heads the maximum- likelihood objective on simulated pairs is

\[ \mathcal{L}_\text{MLE}(\phi) = -\mathbb{E}_{(x, y)} \log q_\phi(x \mid y). \]
Algorithm
  1. Evaluate \(\log q_\phi(x \mid y)\) for each sample with model.log_prob(batch.target, batch).
  2. Return the negative batch mean.

Parameters:

Name Type Description Default
model Any

AmortizedPosterior (any head with a .log_prob method).

 required
batch Batch1D | Batch2D

Training batch with target = x and input = y.

 required

Returns:

Type Description
Float[Array, '']

Scalar NLL averaged over the batch.

Raises:

Type Description
MissingInputError

If batch.target is None.

Examples:

>>> import jax
>>> import jax.numpy as jnp
>>> import optax
>>> import vardax as vdx
>>> key = jax.random.PRNGKey(0)
>>> posterior = vdx.AmortizedPosterior(
...     encoder=vdx.MLPObsEncoder(input_size=12, context_dim=6, key=key),
...     head=vdx.RegressionHead(context_dim=6, state_shape=(3, 4), key=key),
...     config=vdx.AmortizedConfig(),
... )
>>> ones = jnp.ones((2, 3, 4))
>>> batch = vdx.Batch1D(input=ones, mask=ones, target=jnp.zeros((2, 3, 4)))
>>> vdx.amortized_nll_loss_fn(posterior, batch).shape
()

It is the negative mean of the per-sample log-density:

>>> expected = -jnp.mean(posterior.log_prob(batch.target, batch))
>>> bool(jnp.allclose(vdx.amortized_nll_loss_fn(posterior, batch), expected))
True

The true state is required:

>>> blind = vdx.Batch1D(input=ones, mask=ones)
>>> vdx.amortized_nll_loss_fn(
...     posterior, blind
... )
Traceback (most recent call last):
...
MissingInputError: amortized_nll_loss_fn requires batch.target ...
Source code in src/vardax/_src/training.py 
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
def amortized_nll_loss_fn(
    model: Any,
    batch: Batch1D | Batch2D,
) -> Float[Array, ""]:
    r"""Negative log-likelihood for amortized inference (Epic 8).

    For ``AmortizedPosterior`` with flow / regression heads the maximum-
    likelihood objective on simulated pairs is

    $$
    \mathcal{L}_\text{MLE}(\phi) = -\mathbb{E}_{(x, y)} \log q_\phi(x \mid y).
    $$

    Algorithm:
        1. Evaluate $\log q_\phi(x \mid y)$ for each sample with
           ``model.log_prob(batch.target, batch)``.
        2. Return the negative batch mean.

    Args:
        model: ``AmortizedPosterior`` (any head with a ``.log_prob``
            method).
        batch: Training batch with ``target = x`` and ``input = y``.

    Returns:
        Scalar NLL averaged over the batch.

    Raises:
        MissingInputError: If ``batch.target`` is ``None``.

    Examples:
        >>> import jax
        >>> import jax.numpy as jnp
        >>> import optax
        >>> import vardax as vdx
        >>> key = jax.random.PRNGKey(0)
        >>> posterior = vdx.AmortizedPosterior(
        ...     encoder=vdx.MLPObsEncoder(input_size=12, context_dim=6, key=key),
        ...     head=vdx.RegressionHead(context_dim=6, state_shape=(3, 4), key=key),
        ...     config=vdx.AmortizedConfig(),
        ... )
        >>> ones = jnp.ones((2, 3, 4))
        >>> batch = vdx.Batch1D(input=ones, mask=ones, target=jnp.zeros((2, 3, 4)))
        >>> vdx.amortized_nll_loss_fn(posterior, batch).shape
        ()

        It is the negative mean of the per-sample log-density:

        >>> expected = -jnp.mean(posterior.log_prob(batch.target, batch))
        >>> bool(jnp.allclose(vdx.amortized_nll_loss_fn(posterior, batch), expected))
        True

        The true state is required:

        >>> blind = vdx.Batch1D(input=ones, mask=ones)
        >>> vdx.amortized_nll_loss_fn(
        ...     posterior, blind
        ... )  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        MissingInputError: amortized_nll_loss_fn requires batch.target ...
    """
    target = batch.target
    if target is None:
        raise MissingInputError(
            "amortized_nll_loss_fn requires batch.target (the true x) to be set."
        )
    log_prob = model.log_prob(target, batch)
    return -jnp.mean(log_prob)

amortized_train_step 

amortized_train_step(
    model: Any,
    batch: Batch1D | Batch2D,
    optimizer: GradientTransformation,
    opt_state: OptState,
) -> tuple[Any, OptState, Float[Array, ""]]

Single training step for AmortizedPosterior.

Same shape as train_step but uses amortized_nll_loss_fn instead of the MSE reconstruction loss. Use this for simulation-based training of amortized variants; use train_step for FourDVarNet and classical models that reconstruct fields.

Algorithm
  1. Evaluate the NLL and its gradient with respect to the parameters.
  2. Turn the gradient into an update with the optax optimiser.
  3. Apply the update to the model.

Parameters:

Name Type Description Default
model Any

AmortizedPosterior to optimise.

 required
batch Batch1D | Batch2D

Training batch with target = x and input = y.

 required
optimizer GradientTransformation

Optax gradient transformation.

 required
opt_state OptState

Current optimiser state.

 required

Returns:

Type Description
Any

The updated model, the updated optimiser state and the scalar NLL
OptState

(evaluated before the update).

Raises:

Type Description
MissingInputError

If batch.target is None.

Examples:

>>> import jax
>>> import jax.numpy as jnp
>>> import optax
>>> import vardax as vdx
>>> key = jax.random.PRNGKey(0)
>>> posterior = vdx.AmortizedPosterior(
...     encoder=vdx.MLPObsEncoder(input_size=12, context_dim=6, key=key),
...     head=vdx.RegressionHead(context_dim=6, state_shape=(3, 4), key=key),
...     config=vdx.AmortizedConfig(),
... )
>>> ones = jnp.ones((2, 3, 4))
>>> batch = vdx.Batch1D(input=ones, mask=ones, target=jnp.zeros((2, 3, 4)))
>>> import equinox as eqx
>>> optimizer = optax.adam(1e-2)
>>> opt_state = optimizer.init(eqx.filter(posterior, eqx.is_array))
>>> new, opt_state, loss = vdx.amortized_train_step(
...     posterior, batch, optimizer, opt_state
... )
>>> loss.shape
()

A few steps reduce the loss on a fixed batch:

>>> first = loss
>>> for _ in range(5):
...     new, opt_state, loss = vdx.amortized_train_step(
...         new, batch, optimizer, opt_state
...     )
>>> bool(loss < first)
True
Source code in src/vardax/_src/training.py 
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
@eqx.filter_jit
def amortized_train_step(
    model: Any,
    batch: Batch1D | Batch2D,
    optimizer: optax.GradientTransformation,
    opt_state: optax.OptState,
) -> tuple[Any, optax.OptState, Float[Array, ""]]:
    """Single training step for ``AmortizedPosterior``.

    Same shape as [`train_step`][vardax.train_step] but uses
    [`amortized_nll_loss_fn`][vardax.amortized_nll_loss_fn] instead of the MSE
    reconstruction loss. Use this for simulation-based training of amortized
    variants; use ``train_step`` for ``FourDVarNet`` and classical models that
    reconstruct fields.

    Algorithm:
        1. Evaluate the NLL and its gradient with respect to the parameters.
        2. Turn the gradient into an update with the optax optimiser.
        3. Apply the update to the model.

    Args:
        model: ``AmortizedPosterior`` to optimise.
        batch: Training batch with ``target = x`` and ``input = y``.
        optimizer: Optax gradient transformation.
        opt_state: Current optimiser state.

    Returns:
        The updated model, the updated optimiser state and the scalar NLL
        (evaluated before the update).

    Raises:
        MissingInputError: If ``batch.target`` is ``None``.

    Examples:
        >>> import jax
        >>> import jax.numpy as jnp
        >>> import optax
        >>> import vardax as vdx
        >>> key = jax.random.PRNGKey(0)
        >>> posterior = vdx.AmortizedPosterior(
        ...     encoder=vdx.MLPObsEncoder(input_size=12, context_dim=6, key=key),
        ...     head=vdx.RegressionHead(context_dim=6, state_shape=(3, 4), key=key),
        ...     config=vdx.AmortizedConfig(),
        ... )
        >>> ones = jnp.ones((2, 3, 4))
        >>> batch = vdx.Batch1D(input=ones, mask=ones, target=jnp.zeros((2, 3, 4)))
        >>> import equinox as eqx
        >>> optimizer = optax.adam(1e-2)
        >>> opt_state = optimizer.init(eqx.filter(posterior, eqx.is_array))
        >>> new, opt_state, loss = vdx.amortized_train_step(
        ...     posterior, batch, optimizer, opt_state
        ... )
        >>> loss.shape
        ()

        A few steps reduce the loss on a fixed batch:

        >>> first = loss
        >>> for _ in range(5):
        ...     new, opt_state, loss = vdx.amortized_train_step(
        ...         new, batch, optimizer, opt_state
        ...     )
        >>> bool(loss < first)
        True
    """
    loss, grads = eqx.filter_value_and_grad(amortized_nll_loss_fn)(model, batch)
    updates, opt_state = optimizer.update(grads, opt_state, model)
    model = eqx.apply_updates(model, updates)
    return model, opt_state, loss