Amortized Inference

Where the variational methods on the Models page solve an optimisation problem per analysis, an amortized posterior pays the cost once at training time: a network is fit to map observations directly to an approximate posterior, so inference is a single forward pass. See Amortized Inference in the Mathematical Reference for the underlying theory and the fidelity/speed trade-offs.

AmortizedPosterior composes two exchangeable parts: an observation encoder that summarises (possibly masked, possibly irregular) observations into a conditioning vector, and a posterior head that turns that vector into a distribution over states. Heads span the fidelity ladder — point estimates (RegressionHead), full densities via conditional normalizing flows (ConditionalFlowHead), and score-based diffusion sampling (ScoreDiffusionHead). Amortized posteriors should pass the same validation gates (simulation-based calibration, posterior agreement) as their variational counterparts.

Posterior and configuration

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(...)

AmortizedPosterior

Bases: Module

Amortized variational posterior \(q_\phi(x \mid y)\).

Composes an encoder \(c = E(y, m)\) with a density head \(q_\phi(x \mid c)\); a trained model answers any new observation with one forward pass.

Algorithm

1. Encode each batch element's (input, mask) into a context vector (vmapped over the batch).
2. Hand the context to the head: __call__ takes its mode, sample draws from it, log_prob evaluates the density.

Attributes:

Name	Type	Description
encoder	Any	eqx.Module mapping (input, mask) to a context vector. Vmapped over the batch axis at call time.
head	Any	RegressionHead / ConditionalFlowHead / ScoreDiffusionHead. Exposes map_estimate(ctx), sample(ctx, key, n), log_prob(x, ctx).
config	AmortizedConfig	AmortizedConfig carrying head type and sampling defaults.

Examples:

>>> import jax
>>> import jax.numpy as jnp
>>> 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(n_samples=7),
... )
>>> ones = jnp.ones((2, 3, 4))
>>> batch = vdx.Batch1D(input=ones, mask=ones)
>>> posterior(batch).shape
(2, 3, 4)

Samples default to config.n_samples and can be overridden:

>>> posterior.sample(batch, jax.random.PRNGKey(1)).shape
(2, 7, 3, 4)
>>> posterior.sample(batch, jax.random.PRNGKey(1), n=3).shape
(2, 3, 3, 4)

One log-density comes back per batch element, here at the mode:

>>> posterior.log_prob(posterior(batch), batch).shape
(2,)

Source code in src/vardax/_src/amortized/posterior.py

class AmortizedPosterior(eqx.Module):
    r"""Amortized variational posterior $q_\phi(x \mid y)$.

    Composes an encoder $c = E(y, m)$ with a density head $q_\phi(x \mid c)$;
    a trained model answers any new observation with one forward pass.

    Algorithm:
        1. Encode each batch element's ``(input, mask)`` into a context vector
           (vmapped over the batch).
        2. Hand the context to the head: ``__call__`` takes its mode,
           ``sample`` draws from it, ``log_prob`` evaluates the density.

    Attributes:
        encoder: ``eqx.Module`` mapping ``(input, mask)`` to a context
            vector. Vmapped over the batch axis at call time.
        head: ``RegressionHead`` / ``ConditionalFlowHead`` /
            ``ScoreDiffusionHead``. Exposes ``map_estimate(ctx)``,
            ``sample(ctx, key, n)``, ``log_prob(x, ctx)``.
        config: ``AmortizedConfig`` carrying head type and sampling
            defaults.

    Examples:
        >>> import jax
        >>> import jax.numpy as jnp
        >>> 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(n_samples=7),
        ... )
        >>> ones = jnp.ones((2, 3, 4))
        >>> batch = vdx.Batch1D(input=ones, mask=ones)
        >>> posterior(batch).shape
        (2, 3, 4)

        Samples default to ``config.n_samples`` and can be overridden:

        >>> posterior.sample(batch, jax.random.PRNGKey(1)).shape
        (2, 7, 3, 4)
        >>> posterior.sample(batch, jax.random.PRNGKey(1), n=3).shape
        (2, 3, 3, 4)

        One log-density comes back per batch element, here at the mode:

        >>> posterior.log_prob(posterior(batch), batch).shape
        (2,)
    """

    encoder: Any
    head: Any
    config: AmortizedConfig

    def __call__(self, batch: Batch1D | Batch2D) -> Float[Array, ...]:
        r"""Return the per-sample MAP / mode of $q_\phi(x \mid y)$.

        Args:
            batch: Observations ``(input, mask)`` to condition on.

        Returns:
            The mode of the posterior for each batch element, stacked over the batch.
        """

        def _one(input_sample, mask_sample):
            context = self.encoder(input_sample, mask_sample)
            return self.head.map_estimate(context)

        return jax.vmap(_one)(batch.input, batch.mask)

    def sample(
        self,
        batch: Batch1D | Batch2D,
        key: PRNGKeyArray,
        n: int | None = None,
    ) -> Float[Array, ...]:
        """Draw posterior samples per batch element.

        Args:
            batch: Observations ``(input, mask)`` to condition on.
            key: PRNG key, split into one key per batch element.
            n: Number of samples per element. Defaults to ``config.n_samples``.

        Returns:
            Samples stacked as ``(B, n, *state_shape)``.
        """
        n = self.config.n_samples if n is None else n
        batch_size = batch.input.shape[0]
        keys = jax.random.split(key, batch_size)

        def _one(input_sample, mask_sample, key_sample):
            context = self.encoder(input_sample, mask_sample)
            return self.head.sample(context, key_sample, n)

        return jax.vmap(_one)(batch.input, batch.mask, keys)

    def log_prob(
        self,
        x: Float[Array, ...],
        batch: Batch1D | Batch2D,
    ) -> Float[Array, " B"]:
        r"""Per-sample log-density of ``x`` under ``q_φ(·|y)``.

        For ``ScoreDiffusionHead`` this raises ``NotImplementedError``
        (no closed-form density). Use ``sample`` instead.

        Args:
            x: States whose density is evaluated, one per batch element.
            batch: Observations ``(input, mask)`` to condition on.

        Returns:
            The log-density $\log q_\phi(x \mid y)$ for each batch element.
        """

        def _one(state_sample, input_sample, mask_sample):
            context = self.encoder(input_sample, mask_sample)
            return self.head.log_prob(state_sample, context)

        return jax.vmap(_one)(x, batch.input, batch.mask)

    def as_analysis_step(self) -> _AmortizedAnalysisStep:
        """Adapt to ``pipekit_cycle.AnalysisStep`` (Decision D8).

        Returns:
            A callable ``(forecast, obs, *, obs_op, obs_err_cov) -> analysis`` that
            masks the non-finite entries of ``obs`` and returns the MAP estimate.

        Examples:
            Non-finite observations are treated as unobserved:

            >>> import jax
            >>> import jax.numpy as jnp
            >>> 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(),
            ... )
            >>> obs = jnp.ones((1, 3, 4)).at[0, 0, 0].set(jnp.nan)
            >>> step = posterior.as_analysis_step()
            >>> step(obs, obs, obs_op=None, obs_err_cov=None).shape
            (1, 3, 4)
        """
        return _AmortizedAnalysisStep(self)

sample

sample(
    batch: Batch1D | Batch2D,
    key: PRNGKeyArray,
    n: int | None = None,
) -> Float[Array, ...]

Draw posterior samples per batch element.

Parameters:

Name	Type	Description	Default
batch	Batch1D | Batch2D	Observations (input, mask) to condition on.	required
key	PRNGKeyArray	PRNG key, split into one key per batch element.	required
n	int | None	Number of samples per element. Defaults to config.n_samples.	None

Returns:

Type	Description
Float[Array, ...]	Samples stacked as (B, n, *state_shape).

Source code in src/vardax/_src/amortized/posterior.py

def sample(
    self,
    batch: Batch1D | Batch2D,
    key: PRNGKeyArray,
    n: int | None = None,
) -> Float[Array, ...]:
    """Draw posterior samples per batch element.

    Args:
        batch: Observations ``(input, mask)`` to condition on.
        key: PRNG key, split into one key per batch element.
        n: Number of samples per element. Defaults to ``config.n_samples``.

    Returns:
        Samples stacked as ``(B, n, *state_shape)``.
    """
    n = self.config.n_samples if n is None else n
    batch_size = batch.input.shape[0]
    keys = jax.random.split(key, batch_size)

    def _one(input_sample, mask_sample, key_sample):
        context = self.encoder(input_sample, mask_sample)
        return self.head.sample(context, key_sample, n)

    return jax.vmap(_one)(batch.input, batch.mask, keys)

log_prob

log_prob(
    x: Float[Array, ...], batch: Batch1D | Batch2D
) -> Float[Array, " B"]

Per-sample log-density of x under q_φ(·|y).

For ScoreDiffusionHead this raises NotImplementedError (no closed-form density). Use sample instead.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	States whose density is evaluated, one per batch element.	required
batch	Batch1D | Batch2D	Observations (input, mask) to condition on.	required

Returns:

Type	Description
Float[Array, ' B']	The log-density \(\log q_\phi(x \mid y)\) for each batch element.

Source code in src/vardax/_src/amortized/posterior.py

def log_prob(
    self,
    x: Float[Array, ...],
    batch: Batch1D | Batch2D,
) -> Float[Array, " B"]:
    r"""Per-sample log-density of ``x`` under ``q_φ(·|y)``.

    For ``ScoreDiffusionHead`` this raises ``NotImplementedError``
    (no closed-form density). Use ``sample`` instead.

    Args:
        x: States whose density is evaluated, one per batch element.
        batch: Observations ``(input, mask)`` to condition on.

    Returns:
        The log-density $\log q_\phi(x \mid y)$ for each batch element.
    """

    def _one(state_sample, input_sample, mask_sample):
        context = self.encoder(input_sample, mask_sample)
        return self.head.log_prob(state_sample, context)

    return jax.vmap(_one)(x, batch.input, batch.mask)

as_analysis_step

as_analysis_step() -> _AmortizedAnalysisStep

Adapt to pipekit_cycle.AnalysisStep (Decision D8).

Returns:

Type	Description
_AmortizedAnalysisStep	A callable (forecast, obs, *, obs_op, obs_err_cov) -> analysis that
_AmortizedAnalysisStep	masks the non-finite entries of obs and returns the MAP estimate.

Examples:

Non-finite observations are treated as unobserved:

>>> import jax
>>> import jax.numpy as jnp
>>> 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(),
... )
>>> obs = jnp.ones((1, 3, 4)).at[0, 0, 0].set(jnp.nan)
>>> step = posterior.as_analysis_step()
>>> step(obs, obs, obs_op=None, obs_err_cov=None).shape
(1, 3, 4)

Source code in src/vardax/_src/amortized/posterior.py

def as_analysis_step(self) -> _AmortizedAnalysisStep:
    """Adapt to ``pipekit_cycle.AnalysisStep`` (Decision D8).

    Returns:
        A callable ``(forecast, obs, *, obs_op, obs_err_cov) -> analysis`` that
        masks the non-finite entries of ``obs`` and returns the MAP estimate.

    Examples:
        Non-finite observations are treated as unobserved:

        >>> import jax
        >>> import jax.numpy as jnp
        >>> 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(),
        ... )
        >>> obs = jnp.ones((1, 3, 4)).at[0, 0, 0].set(jnp.nan)
        >>> step = posterior.as_analysis_step()
        >>> step(obs, obs, obs_op=None, obs_err_cov=None).shape
        (1, 3, 4)
    """
    return _AmortizedAnalysisStep(self)

AmortizedConfig

Bases: Module

Configuration for an AmortizedPosterior.

Attributes:

Name	Type	Description
head_type	str	One of "regression", "flow", "score". Determines the density family used by the head. Currently only "regression" is fully implemented; "flow" requires gauss_flows and "score" is pending the diffrax reverse-SDE pipeline. Set on construction; the head instance is responsible for matching the type.
n_samples	int	Default number of posterior samples to draw in AmortizedPosterior.sample() when n is not provided.

Raises:

Type	Description
ConfigurationError	If head_type is not one of the three values.

Examples:

The defaults give a regression head and 64 samples:

>>> import vardax as vdx
>>> config = vdx.AmortizedConfig()
>>> config.head_type, config.n_samples
('regression', 64)

Both fields can be set:

>>> vdx.AmortizedConfig(head_type="flow", n_samples=10).n_samples
10

An unknown head type is rejected:

>>> vdx.AmortizedConfig(head_type="bogus")
Traceback (most recent call last):
...
ConfigurationError: AmortizedConfig.head_type must be one of ...

Source code in src/vardax/_src/amortized/config.py

class AmortizedConfig(eqx.Module):
    """Configuration for an ``AmortizedPosterior``.

    Attributes:
        head_type: One of ``"regression"``, ``"flow"``, ``"score"``.
            Determines the density family used by the head. Currently
            only ``"regression"`` is fully implemented;  ``"flow"``
            requires ``gauss_flows`` and ``"score"`` is pending the
            diffrax reverse-SDE pipeline. Set on construction; the head
            instance is responsible for matching the type.
        n_samples: Default number of posterior samples to draw in
            ``AmortizedPosterior.sample()`` when ``n`` is not provided.

    Raises:
        ConfigurationError: If ``head_type`` is not one of the three values.

    Examples:
        The defaults give a regression head and 64 samples:

        >>> import vardax as vdx
        >>> config = vdx.AmortizedConfig()
        >>> config.head_type, config.n_samples
        ('regression', 64)

        Both fields can be set:

        >>> vdx.AmortizedConfig(head_type="flow", n_samples=10).n_samples
        10

        An unknown head type is rejected:

        >>> vdx.AmortizedConfig(head_type="bogus")  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        ConfigurationError: AmortizedConfig.head_type must be one of ...
    """

    head_type: str = eqx.field(static=True, default="regression")
    n_samples: int = eqx.field(static=True, default=64)

    def __post_init__(self) -> None:
        valid = {"regression", "flow", "score"}
        if self.head_type not in valid:
            raise ConfigurationError(
                f"AmortizedConfig.head_type must be one of {sorted(valid)}; "
                f"got {self.head_type!r}."
            )

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

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

Observation encoders

IdentityObsEncoder passes observations straight through — appropriate when they are already a fixed-size vector; MLPObsEncoder learns the summary.

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(...)

IdentityObsEncoder

Bases: Module

Concatenate input and mask into a flat context vector.

Zero parameters, the encoder is purely structural: \(c = [\mathrm{vec}(y); \mathrm{vec}(m)]\). The output dimension is input.size + mask.size.

Examples:

>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> enc = vdx.IdentityObsEncoder()
>>> enc(jnp.ones((3, 4)), jnp.zeros((3, 4))).shape
(24,)

The first half is the data and the second half the mask:

>>> ctx = enc(jnp.ones((2, 2)), jnp.zeros((2, 2)))
>>> ctx.tolist()
[1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0]

It feeds a RegressionHead whose context size is twice the field size:

>>> import jax
>>> head = vdx.RegressionHead(
...     context_dim=24, state_shape=(3, 4), key=jax.random.PRNGKey(0)
... )
>>> head.map_estimate(enc(jnp.ones((3, 4)), jnp.ones((3, 4)))).shape
(3, 4)

Source code in src/vardax/_src/amortized/encoder.py

class IdentityObsEncoder(eqx.Module):
    r"""Concatenate ``input`` and ``mask`` into a flat context vector.

    Zero parameters, the encoder is purely structural:
    $c = [\mathrm{vec}(y); \mathrm{vec}(m)]$. The output dimension is
    ``input.size + mask.size``.

    Examples:
        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> enc = vdx.IdentityObsEncoder()
        >>> enc(jnp.ones((3, 4)), jnp.zeros((3, 4))).shape
        (24,)

        The first half is the data and the second half the mask:

        >>> ctx = enc(jnp.ones((2, 2)), jnp.zeros((2, 2)))
        >>> ctx.tolist()
        [1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0]

        It feeds a [`RegressionHead`][vardax.RegressionHead] whose context size
        is twice the field size:

        >>> import jax
        >>> head = vdx.RegressionHead(
        ...     context_dim=24, state_shape=(3, 4), key=jax.random.PRNGKey(0)
        ... )
        >>> head.map_estimate(enc(jnp.ones((3, 4)), jnp.ones((3, 4)))).shape
        (3, 4)
    """

    def __call__(
        self,
        input: Float[Array, ...],
        mask: Float[Array, ...],
    ) -> Float[Array, " D"]:
        return einx.id("a, b -> (a + b)", flatten(input), flatten(mask))

MLPObsEncoder

Bases: Module

Two-layer MLP encoder from flat (input, mask) to context.

Computes \(c = \mathrm{MLP}([\mathrm{vec}(y); \mathrm{vec}(m)])\).

Parameters:

Name	Type	Description	Default
input_size	int	Flattened size of the input field (T * N for Batch1D, T * H * W for Batch2D).	required
context_dim	int	Dimension of the context vector \(c\).	required
hidden_dim	int	Width of the MLP hidden layers.	64
depth	int	Number of hidden layers of the MLP.	2
key	PRNGKeyArray	PRNG key for the MLP weights.	required

Attributes:

Name	Type	Description
mlp	MLP	eqx.nn.MLP taking 2 * input_size features and emitting context_dim.

Examples:

>>> import jax
>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> enc = vdx.MLPObsEncoder(
...     input_size=12, context_dim=6, key=jax.random.PRNGKey(0)
... )
>>> enc(jnp.ones((3, 4)), jnp.ones((3, 4))).shape
(6,)

Any field shape with the right size is flattened first:

>>> enc(jnp.ones((12,)), jnp.ones((2, 6))).shape
(6,)

It is the encoder of an AmortizedPosterior, where it is vmapped over the batch:

>>> head = vdx.RegressionHead(
...     context_dim=6, state_shape=(3, 4), key=jax.random.PRNGKey(1)
... )
>>> posterior = vdx.AmortizedPosterior(enc, head, vdx.AmortizedConfig())
>>> ones = jnp.ones((2, 3, 4))
>>> posterior(vdx.Batch1D(input=ones, mask=ones)).shape
(2, 3, 4)

Source code in src/vardax/_src/amortized/encoder.py

class MLPObsEncoder(eqx.Module):
    r"""Two-layer MLP encoder from flat ``(input, mask)`` to context.

    Computes $c = \mathrm{MLP}([\mathrm{vec}(y); \mathrm{vec}(m)])$.

    Args:
        input_size: Flattened size of the input field (``T * N`` for
            ``Batch1D``, ``T * H * W`` for ``Batch2D``).
        context_dim: Dimension of the context vector $c$.
        hidden_dim: Width of the MLP hidden layers.
        depth: Number of hidden layers of the MLP.
        key: PRNG key for the MLP weights.

    Attributes:
        mlp: ``eqx.nn.MLP`` taking ``2 * input_size`` features and
            emitting ``context_dim``.

    Examples:
        >>> import jax
        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> enc = vdx.MLPObsEncoder(
        ...     input_size=12, context_dim=6, key=jax.random.PRNGKey(0)
        ... )
        >>> enc(jnp.ones((3, 4)), jnp.ones((3, 4))).shape
        (6,)

        Any field shape with the right size is flattened first:

        >>> enc(jnp.ones((12,)), jnp.ones((2, 6))).shape
        (6,)

        It is the encoder of an [`AmortizedPosterior`][vardax.AmortizedPosterior],
        where it is vmapped over the batch:

        >>> head = vdx.RegressionHead(
        ...     context_dim=6, state_shape=(3, 4), key=jax.random.PRNGKey(1)
        ... )
        >>> posterior = vdx.AmortizedPosterior(enc, head, vdx.AmortizedConfig())
        >>> ones = jnp.ones((2, 3, 4))
        >>> posterior(vdx.Batch1D(input=ones, mask=ones)).shape
        (2, 3, 4)
    """

    mlp: eqx.nn.MLP
    input_size: int = eqx.field(static=True)

    def __init__(
        self,
        input_size: int,
        context_dim: int,
        *,
        hidden_dim: int = 64,
        depth: int = 2,
        key: PRNGKeyArray,
    ) -> None:
        self.input_size = input_size
        self.mlp = eqx.nn.MLP(
            in_size=2 * input_size,
            out_size=context_dim,
            width_size=hidden_dim,
            depth=depth,
            key=key,
        )

    def __call__(
        self,
        input: Float[Array, ...],
        mask: Float[Array, ...],
    ) -> Float[Array, " context_dim"]:
        # features = [vec(y); vec(m)]
        features = einx.id("a, b -> (a + b)", flatten(input), flatten(mask))
        return self.mlp(features)

Posterior heads

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(...)

RegressionHead

Bases: Module

Gaussian regression head: q_φ(x|y) = N(μ_φ(y), diag(σ²_φ(y))).

\[ q_\phi(x \mid y) = \mathcal{N}(\mu_\phi(y), \mathrm{diag}\,\sigma^2_\phi(y)) \]

Two MLPs share the context: one for the mean, one for the log variance. Outputs are reshaped to the user-specified state_shape.

Algorithm

1. Map the context to the flat mean \(\mu\) and log variance \(\log \sigma^2\) with the two MLPs, and reshape them to state_shape.
2. map_estimate returns \(\mu\).
3. sample returns \(\mu + \sigma \odot \varepsilon\) with \(\varepsilon \sim \mathcal{N}(0, I)\).
4. log_prob returns the diagonal-Gaussian log-density summed over all state entries.

Parameters:

Name	Type	Description	Default
context_dim	int	Dimension of the context vector from the encoder.	required
state_shape	tuple[int, ...]	Shape of a single posterior sample (e.g. (T, N) for Batch1D targets).	required
hidden_dim	int	Width of the MLP hidden layers.	64
depth	int	Number of hidden layers of each MLP.	2
key	PRNGKeyArray	PRNG key for the MLP weights.	required

Attributes:

Name	Type	Description
mlp_mu	MLP	MLP from context to flat mean (size prod(state_shape)).
mlp_log_var	MLP	MLP from context to flat log variance.

Examples:

>>> import jax
>>> import jax.numpy as jnp
>>> import vardax as vdx
>>> head = vdx.RegressionHead(
...     context_dim=6, state_shape=(3, 4), key=jax.random.PRNGKey(0)
... )
>>> ctx = jnp.ones(6)
>>> head.map_estimate(ctx).shape
(3, 4)
>>> head.sample(ctx, jax.random.PRNGKey(1), 5).shape
(5, 3, 4)

The density is highest at the mean:

>>> mode = head.map_estimate(ctx)
>>> bool(head.log_prob(mode, ctx) > head.log_prob(mode + 3.0, ctx))
True

The head scalar log-density has no batch axis:

>>> head.log_prob(mode, ctx).shape
()

Source code in src/vardax/_src/amortized/heads.py

class RegressionHead(eqx.Module):
    r"""Gaussian regression head: ``q_φ(x|y) = N(μ_φ(y), diag(σ²_φ(y)))``.

    $$
    q_\phi(x \mid y) = \mathcal{N}(\mu_\phi(y), \mathrm{diag}\,\sigma^2_\phi(y))
    $$

    Two MLPs share the context: one for the mean, one for the log
    variance. Outputs are reshaped to the user-specified
    ``state_shape``.

    Algorithm:
        1. Map the context to the flat mean $\mu$ and log variance
           $\log \sigma^2$ with the two MLPs, and reshape them to ``state_shape``.
        2. ``map_estimate`` returns $\mu$.
        3. ``sample`` returns $\mu + \sigma \odot \varepsilon$ with
           $\varepsilon \sim \mathcal{N}(0, I)$.
        4. ``log_prob`` returns the diagonal-Gaussian log-density summed over
           all state entries.

    Args:
        context_dim: Dimension of the context vector from the encoder.
        state_shape: Shape of a single posterior sample (e.g.
            ``(T, N)`` for ``Batch1D`` targets).
        hidden_dim: Width of the MLP hidden layers.
        depth: Number of hidden layers of each MLP.
        key: PRNG key for the MLP weights.

    Attributes:
        mlp_mu: MLP from context to flat mean (size ``prod(state_shape)``).
        mlp_log_var: MLP from context to flat log variance.

    Examples:
        >>> import jax
        >>> import jax.numpy as jnp
        >>> import vardax as vdx
        >>> head = vdx.RegressionHead(
        ...     context_dim=6, state_shape=(3, 4), key=jax.random.PRNGKey(0)
        ... )
        >>> ctx = jnp.ones(6)
        >>> head.map_estimate(ctx).shape
        (3, 4)
        >>> head.sample(ctx, jax.random.PRNGKey(1), 5).shape
        (5, 3, 4)

        The density is highest at the mean:

        >>> mode = head.map_estimate(ctx)
        >>> bool(head.log_prob(mode, ctx) > head.log_prob(mode + 3.0, ctx))
        True

        The head scalar log-density has no batch axis:

        >>> head.log_prob(mode, ctx).shape
        ()
    """

    mlp_mu: eqx.nn.MLP
    mlp_log_var: eqx.nn.MLP
    state_shape: tuple[int, ...] = eqx.field(static=True)

    def __init__(
        self,
        context_dim: int,
        state_shape: tuple[int, ...],
        *,
        hidden_dim: int = 64,
        depth: int = 2,
        key: PRNGKeyArray,
    ) -> None:
        self.state_shape = tuple(state_shape)
        flat_dim = 1
        for dim_size in self.state_shape:
            flat_dim *= dim_size
        key_mean, key_log_var = jax.random.split(key)
        self.mlp_mu = eqx.nn.MLP(
            in_size=context_dim,
            out_size=flat_dim,
            width_size=hidden_dim,
            depth=depth,
            key=key_mean,
        )
        self.mlp_log_var = eqx.nn.MLP(
            in_size=context_dim,
            out_size=flat_dim,
            width_size=hidden_dim,
            depth=depth,
            key=key_log_var,
        )

    def map_estimate(self, ctx: Float[Array, " D"]) -> Float[Array, ...]:
        return einx.id("(s...) -> s...", self.mlp_mu(ctx), s=self.state_shape)

    def sample(
        self,
        ctx: Float[Array, " D"],
        key: PRNGKeyArray,
        n: int,
    ) -> Float[Array, ...]:
        mean = einx.id("(s...) -> s...", self.mlp_mu(ctx), s=self.state_shape)
        log_var = einx.id("(s...) -> s...", self.mlp_log_var(ctx), s=self.state_shape)
        # σ = exp(½ log σ²)
        std = jnp.exp(0.5 * log_var)
        noise = jax.random.normal(key, (n, *self.state_shape))
        # x = μ + σ ⊙ ε
        return mean + std * noise

    def log_prob(
        self,
        x: Float[Array, ...],
        ctx: Float[Array, " D"],
    ) -> Float[Array, ""]:
        mean = einx.id("(s...) -> s...", self.mlp_mu(ctx), s=self.state_shape)
        log_var = einx.id("(s...) -> s...", self.mlp_log_var(ctx), s=self.state_shape)
        # log N(x | mu, diag(exp(log_var))) — sum over state dims.
        return -0.5 * jnp.sum(
            ((x - mean) ** 2) * jnp.exp(-log_var) + log_var + _LOG_2PI
        )

ConditionalFlowHead

Bases: Module

Conditional normalising flow head (stub).

Implements \(x = f_\phi(z; c_\psi(y))\) with \(z \sim \mathcal{N}(0, I)\). Exact density via change-of-variables. Requires gauss_flows for the flow primitives; ships as a stub until that dependency is added to pyproject.toml.

Examples:

Constructing the stub raises until the dependency lands:

>>> import vardax as vdx
>>> vdx.ConditionalFlowHead()
Traceback (most recent call last):
...
NotImplementedError: ConditionalFlowHead requires `gauss_flows` ...

Source code in src/vardax/_src/amortized/heads.py

class ConditionalFlowHead(eqx.Module):
    r"""Conditional normalising flow head (stub).

    Implements $x = f_\phi(z; c_\psi(y))$ with $z \sim \mathcal{N}(0, I)$. Exact
    density via change-of-variables. Requires ``gauss_flows`` for the
    flow primitives; ships as a stub until that dependency is added to
    ``pyproject.toml``.

    Examples:
        Constructing the stub raises until the dependency lands:

        >>> import vardax as vdx
        >>> vdx.ConditionalFlowHead()  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        NotImplementedError: ConditionalFlowHead requires `gauss_flows` ...
    """

    def __init__(self, *_args: object, **_kwargs: object) -> None:
        raise NotImplementedError(
            "ConditionalFlowHead requires `gauss_flows`, which is not yet "
            "a vardax dependency. Use `RegressionHead` for now; flow support "
            "lands in a follow-up PR alongside the gauss_flows dependency."
        )

ScoreDiffusionHead

Bases: Module

Score-based diffusion head (stub).

Learns \(s_\phi(x, t \mid y) \approx \nabla_x \log p_t(x \mid y)\); samples via reverse SDE solved with diffrax. Ships as a stub pending the reverse-SDE pipeline.

Examples:

Constructing the stub raises until the pipeline lands:

>>> import vardax as vdx
>>> vdx.ScoreDiffusionHead()
Traceback (most recent call last):
...
NotImplementedError: ScoreDiffusionHead is not yet implemented. ...

Source code in src/vardax/_src/amortized/heads.py

class ScoreDiffusionHead(eqx.Module):
    r"""Score-based diffusion head (stub).

    Learns $s_\phi(x, t \mid y) \approx \nabla_x \log p_t(x \mid y)$; samples via
    reverse SDE solved with diffrax. Ships as a stub pending the reverse-SDE
    pipeline.

    Examples:
        Constructing the stub raises until the pipeline lands:

        >>> import vardax as vdx
        >>> vdx.ScoreDiffusionHead()  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        NotImplementedError: ScoreDiffusionHead is not yet implemented. ...
    """

    def __init__(self, *_args: object, **_kwargs: object) -> None:
        raise NotImplementedError(
            "ScoreDiffusionHead is not yet implemented. Pending the diffrax-based "
            "reverse-SDE pipeline. Use `RegressionHead` for now."
        )