stagereg.jax

Every symbol exported from stagereg.jax, the JAX/Flax NNX training backend, grouped by the module that defines it. Using it needs the jax extra.

stagereg.jax.bayes_by_backprop

BayesByBackprop dataclass

BayesByBackprop(
    optimizer: GradientTransformation = _default_optimizer(),
    initial_std: float = 0.001,
)

Mean-field Gaussian variational inference by Bayes by Backprop.

Every update takes one reparameterized parameter sample and steps the variational parameters along the gradient of the negative minibatch ELBO: the KL divergence from the prior minus the problem's log-likelihood, which is already scaled by N / B. The KL is in closed form for an IsotropicGaussianPrior and a one-sample Monte Carlo estimate, log q(w) - log p(w), for any other prior. Standard deviations are parameterized as the softplus of an unconstrained rho.

The variational means are warm-started from the problem's initial parameters, the trained mean network, and every standard deviation from initial_std. Its default of 1e-3 keeps the first draws close to the trained mean, so inference starts from a good fit and widens the posterior from there; raise it when the parameters' posterior spread is expected to be far larger, since a softplus scale grows only about linearly in rho.

It supports early stopping. Its best checkpoint holds the variational parameters and optimizer state; its finalized checkpoint holds the variational means and standard deviations alone.

Attributes:

Name	Type	Description
optimizer	GradientTransformation	The Optax transformation stepping the variational parameters. Defaults to Adam with a learning rate of 1e-3.
initial_std	float	The standard deviation every parameter starts with, positive and finite. Defaults to 1e-3.

supports_early_stopping property

supports_early_stopping: bool

Always True: the variational parameters can stop anywhere.

has_posterior property

has_posterior: bool

Whether the method has been initialized.

init

init(problem: PosteriorProblem) -> None

Warm-start the variational parameters from problem.

Parameters:

Name	Type	Description	Default
problem	PosteriorProblem	The posterior problem.	required

Source code in src/stagereg/jax/bayes_by_backprop.py

def init(self, problem: PosteriorProblem) -> None:
    """Warm-start the variational parameters from ``problem``.

    Args:
        problem: The posterior problem.
    """
    rho = _inverse_softplus(self.initial_std)
    self._problem = problem
    self._variational = {
        "means": jax.tree.map(jnp.copy, problem.initial_parameters),
        "rhos": jax.tree.map(
            lambda leaf: jnp.full_like(leaf, rho), problem.initial_parameters
        ),
    }
    self._optimizer_state = self.optimizer.init(self._variational)
    self._step = _make_step(problem, self.optimizer)

update

update(batch: Batch, key: Array) -> Mapping[str, Any]

Step the variational parameters on one training batch.

Only the batch's inputs, targets and sample_weight reach the log-likelihood; its metadata does not.

Parameters:

Name	Type	Description	Default
batch	Batch	The training batch.	required
key	Array	A fresh PRNG key for this batch's parameter sample.	required

Returns:

Type	Description
Mapping[str, Any]	{"loss": negative_elbo}, the minibatch estimate of the
Mapping[str, Any]	negative ELBO.

Raises:

Type	Description
ValueError	If the method has not been initialized.

Source code in src/stagereg/jax/bayes_by_backprop.py

def update(self, batch: Batch, key: jax.Array) -> Mapping[str, Any]:
    """Step the variational parameters on one training batch.

    Only the batch's ``inputs``, ``targets`` and ``sample_weight`` reach
    the log-likelihood; its ``metadata`` does not.

    Args:
        batch: The training batch.
        key: A fresh PRNG key for this batch's parameter sample.

    Returns:
        ``{"loss": negative_elbo}``, the minibatch estimate of the
        negative ELBO.

    Raises:
        ValueError: If the method has not been initialized.
    """
    self._require_problem()
    assert self._step is not None
    self._variational, self._optimizer_state, loss = self._step(
        self._variational,
        self._optimizer_state,
        batch.inputs,
        batch.targets,
        batch.sample_weight,
        key,
    )
    return {"loss": loss}

posterior

posterior() -> MeanFieldGaussianPosterior

Return the current variational posterior.

Returns:

Type	Description
MeanFieldGaussianPosterior	A posterior independent of later updates.

Raises:

Type	Description
ValueError	If the method has not been initialized.

Source code in src/stagereg/jax/bayes_by_backprop.py

def posterior(self) -> MeanFieldGaussianPosterior:
    """Return the current variational posterior.

    Returns:
        A posterior independent of later updates.

    Raises:
        ValueError: If the method has not been initialized.
    """
    problem = self._require_problem()
    return MeanFieldGaussianPosterior(
        mean_function=problem.mean_function,
        means=self._variational["means"],
        stds=jax.tree.map(jax.nn.softplus, self._variational["rhos"]),
    )

state

state() -> PyTree

Return the variational parameters and the optimizer state.

Returns:

Type	Description
PyTree	{"variational": {"means", "rhos"}, "optimizer_state"}.

Raises:

Type	Description
ValueError	If the method has not been initialized.

Source code in src/stagereg/jax/bayes_by_backprop.py

def state(self) -> PyTree:
    """Return the variational parameters and the optimizer state.

    Returns:
        ``{"variational": {"means", "rhos"}, "optimizer_state"}``.

    Raises:
        ValueError: If the method has not been initialized.
    """
    self._require_problem()
    return {
        "variational": self._variational,
        "optimizer_state": self._optimizer_state,
    }

load_state

load_state(state: PyTree) -> None

Resume inference from a state returned by state().

Parameters:

Name	Type	Description	Default
state	PyTree	A state previously returned by state, after init on the same problem.	required

Raises:

Type	Description
ValueError	If the method has not been initialized.

Source code in src/stagereg/jax/bayes_by_backprop.py

def load_state(self, state: PyTree) -> None:
    """Resume inference from a state returned by ``state()``.

    Args:
        state: A state previously returned by
            [`state`][stagereg.jax.BayesByBackprop.state], after
            [`init`][stagereg.jax.BayesByBackprop.init] on the same
            problem.

    Raises:
        ValueError: If the method has not been initialized.
    """
    self._require_problem()
    self._variational = state["variational"]
    self._optimizer_state = state["optimizer_state"]

posterior_state

posterior_state() -> PyTree

Return the variational means and standard deviations.

Returns:

Type	Description
PyTree	{"means", "stds"}, each in the posterior network's parameter
PyTree	tree.

Raises:

Type	Description
ValueError	If the method has not been initialized.

Source code in src/stagereg/jax/bayes_by_backprop.py

def posterior_state(self) -> PyTree:
    """Return the variational means and standard deviations.

    Returns:
        ``{"means", "stds"}``, each in the posterior network's parameter
        tree.

    Raises:
        ValueError: If the method has not been initialized.
    """
    posterior = self.posterior()
    return {"means": posterior.means, "stds": posterior.stds}

load_posterior

load_posterior(state: PyTree) -> None

Rebuild the posterior from a state returned by posterior_state().

Parameters:

Name	Type	Description	Default
state	PyTree	A state previously returned by posterior_state.	required

Raises:

Type	Description
ValueError	If the method has not been initialized.

Source code in src/stagereg/jax/bayes_by_backprop.py

def load_posterior(self, state: PyTree) -> None:
    """Rebuild the posterior from a state returned by ``posterior_state()``.

    Args:
        state: A state previously returned by
            [`posterior_state`][stagereg.jax.BayesByBackprop.posterior_state].

    Raises:
        ValueError: If the method has not been initialized.
    """
    self._require_problem()
    self._variational = {
        "means": state["means"],
        "rhos": jax.tree.map(_inverse_softplus, state["stds"]),
    }

MeanFieldGaussianPosterior dataclass

MeanFieldGaussianPosterior(
    mean_function: Callable[[PyTree, PyTree], Array],
    means: PyTree,
    stds: PyTree,
)

An independent Gaussian on every parameter of the posterior network.

Draw s under key samples its parameters from jax.random.fold_in(key, s), so it is the same mean function at every input, in every call, and whatever the number of draws asked for.

Attributes:

Name	Type	Description
mean_function	Callable[[PyTree, PyTree], Array]	Maps (parameters, inputs) to the posterior network's mean predictions.
means	PyTree	The variational means, in the posterior network's parameter tree.
stds	PyTree	The variational standard deviations, positive, in the same tree.

num_draws property

num_draws: int | None

Always None: a variational posterior has unlimited draws.

sample_means

sample_means(
    inputs: PyTree,
    key: Array,
    num_samples: int | None = None,
) -> jax.Array

Draw mean predictions at inputs.

Parameters:

Name	Type	Description	Default
inputs	PyTree	Model inputs with a leading batch dimension.	required
key	Array	The PRNG key selecting the draws.	required
num_samples	int | None	The number of draws S, positive.	None

Returns:

Type	Description
Array	The draws' mean predictions, shaped [S, *batch].

Raises:

Type	Description
ValueError	If num_samples is omitted or not positive.

Source code in src/stagereg/jax/bayes_by_backprop.py

def sample_means(
    self,
    inputs: PyTree,
    key: jax.Array,
    num_samples: int | None = None,
) -> jax.Array:
    """Draw mean predictions at ``inputs``.

    Args:
        inputs: Model inputs with a leading batch dimension.
        key: The PRNG key selecting the draws.
        num_samples: The number of draws ``S``, positive.

    Returns:
        The draws' mean predictions, shaped ``[S, *batch]``.

    Raises:
        ValueError: If ``num_samples`` is omitted or not positive.
    """
    if num_samples is None:
        raise ValueError("an unlimited posterior needs num_samples.")
    if num_samples <= 0:
        raise ValueError("num_samples must be positive.")

    def draw(index: jax.Array) -> jax.Array:
        parameters = _sample(self.means, self.stds, jax.random.fold_in(key, index))
        return self.mean_function(parameters, inputs)

    return jax.vmap(draw)(jnp.arange(num_samples))

stagereg.jax.distributions

Gaussian dataclass

Gaussian(loc: Array, scale: Array)

A JAX-backed Gaussian predictive distribution parametrized by scale.

Attributes:

Name	Type	Description
loc	Array	The distribution mean, shaped like the model's output.
scale	Array	The distribution's positive standard deviation, broadcastable against loc. Callers (e.g. GaussianHead) are responsible for ensuring positivity, e.g. via a softplus transform of an unconstrained raw output.

batch_shape property

batch_shape: tuple[int, ...]

The broadcast shape of independent Gaussian components.

event_shape property

event_shape: tuple[int, ...]

The shape of a single Gaussian event, always scalar.

log_prob

log_prob(targets: Array) -> jax.Array

Compute the elementwise Gaussian log-density of targets.

Parameters:

Name	Type	Description	Default
targets	Array	Target values broadcastable against loc/scale.	required

Returns:

Type	Description
Array	The elementwise log-density, shaped like the broadcast of
Array	targets against loc and scale.

Source code in src/stagereg/jax/distributions.py

def log_prob(self, targets: jax.Array) -> jax.Array:
    """Compute the elementwise Gaussian log-density of ``targets``.

    Args:
        targets: Target values broadcastable against ``loc``/``scale``.

    Returns:
        The elementwise log-density, shaped like the broadcast of
        ``targets`` against ``loc`` and ``scale``.
    """
    return jstats.norm.logpdf(targets, loc=self.loc, scale=self.scale)

sample

sample(
    key: Array, sample_shape: tuple[int, ...] = ()
) -> jax.Array

Draw reparametrized samples from this distribution.

Parameters:

Name	Type	Description	Default
key	Array	A JAX PRNG key.	required
sample_shape	tuple[int, ...]	Leading sample dimensions prepended to batch_shape.	()

Returns:

Type	Description
Array	Samples shaped sample_shape + batch_shape.

Source code in src/stagereg/jax/distributions.py

def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> jax.Array:
    """Draw reparametrized samples from this distribution.

    Args:
        key: A JAX PRNG key.
        sample_shape: Leading sample dimensions prepended to
            ``batch_shape``.

    Returns:
        Samples shaped ``sample_shape + batch_shape``.
    """
    shape = sample_shape + self.batch_shape
    noise = jax.random.normal(key, shape)
    return self.loc + self.scale * noise

mean

mean() -> jax.Array

Return the distribution mean, i.e. loc.

Source code in src/stagereg/jax/distributions.py

def mean(self) -> jax.Array:
    """Return the distribution mean, i.e. ``loc``."""
    return self.loc

variance

variance() -> jax.Array

Return the elementwise variance, i.e. scale ** 2.

Source code in src/stagereg/jax/distributions.py

def variance(self) -> jax.Array:
    """Return the elementwise variance, i.e. ``scale ** 2``."""
    return jnp.square(self.scale)

Gamma dataclass

Gamma(concentration: Array, rate: Array)

A JAX-backed Gamma distribution parametrized by shape and rate.

Attributes:

Name	Type	Description
concentration	Array	Positive Gamma shape parameter.
rate	Array	Positive Gamma rate parameter, i.e. the inverse scale.

batch_shape property

batch_shape: tuple[int, ...]

Return the broadcast shape of independent Gamma components.

event_shape property

event_shape: tuple[int, ...]

Return the shape of a single Gamma event, always scalar.

log_prob

log_prob(targets: Array) -> jax.Array

Compute the elementwise Gamma log-density of targets.

Parameters:

Name	Type	Description	Default
targets	Array	Positive target values broadcastable against the concentration and rate.	required

Returns:

Type	Description
Array	Elementwise log-density values.

Source code in src/stagereg/jax/distributions.py

def log_prob(self, targets: jax.Array) -> jax.Array:
    """Compute the elementwise Gamma log-density of ``targets``.

    Args:
        targets: Positive target values broadcastable against the
            concentration and rate.

    Returns:
        Elementwise log-density values.
    """
    return (
        self.concentration * jnp.log(self.rate)
        - jsp.gammaln(self.concentration)
        + (self.concentration - 1.0) * jnp.log(targets)
        - self.rate * targets
    )

sample

sample(
    key: Array, sample_shape: tuple[int, ...] = ()
) -> jax.Array

Draw keyed samples from the Gamma distribution.

Parameters:

Name	Type	Description	Default
key	Array	A JAX PRNG key.	required
sample_shape	tuple[int, ...]	Leading sample dimensions prepended to batch_shape.	()

Returns:

Type	Description
Array	Samples shaped sample_shape + batch_shape.

Source code in src/stagereg/jax/distributions.py

def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> jax.Array:
    """Draw keyed samples from the Gamma distribution.

    Args:
        key: A JAX PRNG key.
        sample_shape: Leading sample dimensions prepended to
            ``batch_shape``.

    Returns:
        Samples shaped ``sample_shape + batch_shape``.
    """
    shape = sample_shape + self.batch_shape
    concentration = jnp.broadcast_to(self.concentration, self.batch_shape)
    rate = jnp.broadcast_to(self.rate, self.batch_shape)
    return jax.random.gamma(key, concentration, shape=shape) / rate

mean

mean() -> jax.Array

Return the Gamma mean, i.e. concentration / rate.

Source code in src/stagereg/jax/distributions.py

def mean(self) -> jax.Array:
    """Return the Gamma mean, i.e. ``concentration / rate``."""
    return self.concentration / self.rate

variance

variance() -> jax.Array

Return the Gamma variance, i.e. concentration / rate ** 2.

Source code in src/stagereg/jax/distributions.py

def variance(self) -> jax.Array:
    """Return the Gamma variance, i.e. ``concentration / rate ** 2``."""
    return self.concentration / jnp.square(self.rate)

PosteriorPredictive dataclass

PosteriorPredictive(
    draws: Array, aleatoric_variance: Array
)

The posterior predictive: an equally weighted mixture of Gaussians over draws.

Each component is a Gaussian centred on one draw's mean with the aleatoric variance, so the mixture is exact for the draws taken and may be skewed or multimodal. Its single-Gaussian summary is moment_matched.

Attributes:

Name	Type	Description
draws	Array	The mean of each draw at the predicted inputs, shaped [S, *batch] with S >= 1 draws on the leading axis.
aleatoric_variance	Array	The positive aleatoric variance, broadcastable against draws[0] and shared by every draw.

num_draws property

num_draws: int

The number of draws S the mixture averages over.

batch_shape property

batch_shape: tuple[int, ...]

The broadcast shape of the draws' means and the aleatoric variance.

event_shape property

event_shape: tuple[int, ...]

The shape of a single predictive sample, always scalar.

log_prob

log_prob(targets: Array) -> jax.Array

Compute the exact elementwise mixture log-density of targets.

Parameters:

Name	Type	Description	Default
targets	Array	Target values broadcastable against batch_shape.	required

Returns:

Type	Description
Array	logsumexp over draws of the component log-densities, minus
Array	log S.

Source code in src/stagereg/jax/distributions.py

def log_prob(self, targets: jax.Array) -> jax.Array:
    """Compute the exact elementwise mixture log-density of ``targets``.

    Args:
        targets: Target values broadcastable against ``batch_shape``.

    Returns:
        ``logsumexp`` over draws of the component log-densities, minus
        ``log S``.
    """
    loc, scale = self._components()
    component = jstats.norm.logpdf(targets, loc=loc, scale=scale)
    return jsp.logsumexp(component, axis=0) - jnp.log(self.num_draws)

sample

sample(
    key: Array, sample_shape: tuple[int, ...] = ()
) -> jax.Array

Draw predictive samples, each from a uniformly chosen draw's Gaussian.

The Gaussian noise uses key exactly as Gaussian.sample does, so with one draw the samples equal that Gaussian's.

Parameters:

Name	Type	Description	Default
key	Array	A JAX PRNG key.	required
sample_shape	tuple[int, ...]	Leading sample dimensions prepended to batch_shape.	()

Returns:

Type	Description
Array	Predictive samples shaped sample_shape + batch_shape.

Source code in src/stagereg/jax/distributions.py

def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> jax.Array:
    """Draw predictive samples, each from a uniformly chosen draw's Gaussian.

    The Gaussian noise uses ``key`` exactly as
    [`Gaussian.sample`][stagereg.jax.Gaussian.sample] does, so with one draw
    the samples equal that Gaussian's.

    Args:
        key: A JAX PRNG key.
        sample_shape: Leading sample dimensions prepended to
            ``batch_shape``.

    Returns:
        Predictive samples shaped ``sample_shape + batch_shape``.
    """
    shape = sample_shape + self.batch_shape
    noise = jax.random.normal(key, shape)
    component = jax.random.randint(
        jax.random.fold_in(key, 1), shape, 0, self.num_draws
    )
    means = jnp.broadcast_to(
        jnp.moveaxis(self.draws, 0, -1), shape + (self.num_draws,)
    )
    loc = jnp.take_along_axis(means, component[..., None], axis=-1)[..., 0]
    return loc + jnp.sqrt(self.aleatoric_variance) * noise

mean

mean() -> jax.Array

Return the mixture mean, the average of the draws' means.

Source code in src/stagereg/jax/distributions.py

def mean(self) -> jax.Array:
    """Return the mixture mean, the average of the draws' means."""
    return jnp.broadcast_to(jnp.mean(self.draws, axis=0), self.batch_shape)

variance

variance() -> jax.Array

Return the mixture variance, aleatoric plus epistemic.

Source code in src/stagereg/jax/distributions.py

def variance(self) -> jax.Array:
    """Return the mixture variance, aleatoric plus epistemic."""
    return self.moment_matched().variance()

cdf

cdf(targets: Array) -> jax.Array

Compute the elementwise mixture cumulative distribution of targets.

Parameters:

Name	Type	Description	Default
targets	Array	Target values broadcastable against batch_shape.	required

Returns:

Type	Description
Array	The average over draws of the component Gaussian CDFs.

Source code in src/stagereg/jax/distributions.py

def cdf(self, targets: jax.Array) -> jax.Array:
    """Compute the elementwise mixture cumulative distribution of ``targets``.

    Args:
        targets: Target values broadcastable against ``batch_shape``.

    Returns:
        The average over draws of the component Gaussian CDFs.
    """
    loc, scale = self._components()
    component = jstats.norm.cdf(targets, loc=loc, scale=scale)
    return jnp.mean(component, axis=0)

quantile

quantile(probability: float) -> jax.Array

Compute the elementwise mixture quantile at probability.

The mixture CDF has no closed-form inverse, so the quantile is found by bisection between the smallest and the largest of the draws' Gaussian quantiles, which bracket it.

Parameters:

Name	Type	Description	Default
probability	float	The cumulative probability, strictly between 0 and 1.	required

Returns:

Type	Description
Array	The quantile, shaped like batch_shape.

Raises:

Type	Description
ValueError	If probability is not strictly between 0 and 1.

Source code in src/stagereg/jax/distributions.py

def quantile(self, probability: float) -> jax.Array:
    """Compute the elementwise mixture quantile at ``probability``.

    The mixture CDF has no closed-form inverse, so the quantile is found
    by bisection between the smallest and the largest of the draws'
    Gaussian quantiles, which bracket it.

    Args:
        probability: The cumulative probability, strictly between 0 and 1.

    Returns:
        The quantile, shaped like ``batch_shape``.

    Raises:
        ValueError: If ``probability`` is not strictly between 0 and 1.
    """
    if not 0.0 < probability < 1.0:
        raise ValueError("probability must be strictly between 0 and 1.")
    loc, scale = self._components()
    component = loc + NormalDist().inv_cdf(probability) * scale
    lower = jnp.min(component, axis=0)
    upper = jnp.max(component, axis=0)

    def bisect(
        _: int, bracket: tuple[jax.Array, jax.Array]
    ) -> tuple[jax.Array, jax.Array]:
        low, high = bracket
        middle = 0.5 * (low + high)
        below = self.cdf(middle) < probability
        return jnp.where(below, middle, low), jnp.where(below, high, middle)

    lower, upper = jax.lax.fori_loop(0, _BISECTION_STEPS, bisect, (lower, upper))
    return 0.5 * (lower + upper)

moment_matched

moment_matched() -> MomentMatchedPredictive

Summarize the mixture by the Gaussian with its mean and variance.

By the law of total variance, the mixture variance is the aleatoric variance plus the epistemic variance of the draws' means.

Returns:

Type	Description
MomentMatchedPredictive	The MomentMatchedPredictive,
MomentMatchedPredictive	with both variance parts broadcast to batch_shape.

Source code in src/stagereg/jax/distributions.py

def moment_matched(self) -> MomentMatchedPredictive:
    """Summarize the mixture by the Gaussian with its mean and variance.

    By the law of total variance, the mixture variance is the aleatoric
    variance plus the epistemic variance of the draws' means.

    Returns:
        The [`MomentMatchedPredictive`][stagereg.jax.MomentMatchedPredictive],
        with both variance parts broadcast to ``batch_shape``.
    """
    return MomentMatchedPredictive(
        loc=self.mean(),
        aleatoric_variance=jnp.broadcast_to(
            self.aleatoric_variance, self.batch_shape
        ),
        epistemic_variance=jnp.broadcast_to(
            jnp.var(self.draws, axis=0), self.batch_shape
        ),
    )

MomentMatchedPredictive dataclass

MomentMatchedPredictive(
    loc: Array,
    aleatoric_variance: Array,
    epistemic_variance: Array,
)

The Gaussian with a posterior predictive's mean and variance.

A named summary of a PosteriorPredictive, not a substitute for it: its variance is split into the aleatoric and the epistemic part. Obtain one from PosteriorPredictive.moment_matched.

Attributes:

Name	Type	Description
loc	Array	The predictive mean, the average of the draws' means.
aleatoric_variance	Array	The aleatoric variance shared by every draw.
epistemic_variance	Array	The variance of the draws' means around loc.

batch_shape property

batch_shape: tuple[int, ...]

The broadcast shape of the mean and both variance parts.

event_shape property

event_shape: tuple[int, ...]

The shape of a single predictive sample, always scalar.

log_prob

log_prob(targets: Array) -> jax.Array

Compute the elementwise Gaussian log-density of targets.

Parameters:

Name	Type	Description	Default
targets	Array	Target values broadcastable against batch_shape.	required

Returns:

Type	Description
Array	The elementwise log-density under the moment-matched Gaussian.

Source code in src/stagereg/jax/distributions.py

def log_prob(self, targets: jax.Array) -> jax.Array:
    """Compute the elementwise Gaussian log-density of ``targets``.

    Args:
        targets: Target values broadcastable against ``batch_shape``.

    Returns:
        The elementwise log-density under the moment-matched Gaussian.
    """
    return self._gaussian().log_prob(targets)

sample

sample(
    key: Array, sample_shape: tuple[int, ...] = ()
) -> jax.Array

Draw predictive samples from the moment-matched Gaussian.

Parameters:

Name	Type	Description	Default
key	Array	A JAX PRNG key.	required
sample_shape	tuple[int, ...]	Leading sample dimensions prepended to batch_shape.	()

Returns:

Type	Description
Array	Predictive samples shaped sample_shape + batch_shape.

Source code in src/stagereg/jax/distributions.py

def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> jax.Array:
    """Draw predictive samples from the moment-matched Gaussian.

    Args:
        key: A JAX PRNG key.
        sample_shape: Leading sample dimensions prepended to
            ``batch_shape``.

    Returns:
        Predictive samples shaped ``sample_shape + batch_shape``.
    """
    return jnp.broadcast_to(
        self._gaussian().sample(key, sample_shape),
        sample_shape + self.batch_shape,
    )

mean

mean() -> jax.Array

Return the predictive mean, i.e. loc.

Source code in src/stagereg/jax/distributions.py

def mean(self) -> jax.Array:
    """Return the predictive mean, i.e. ``loc``."""
    return self.loc

variance

variance() -> jax.Array

Return the aleatoric plus the epistemic variance.

Source code in src/stagereg/jax/distributions.py

def variance(self) -> jax.Array:
    """Return the aleatoric plus the epistemic variance."""
    return self.aleatoric_variance + self.epistemic_variance

GaussianHead

GaussianHead(
    in_features: int,
    out_features: int,
    *,
    rngs: Rngs,
    eps: float = 1e-06,
)

Bases: Module

An NNX head producing a Gaussian from model features.

A single linear layer maps in_features to 2 * out_features outputs, split into an unconstrained location and an unconstrained scale that is passed through softplus (plus eps) to guarantee a strictly positive standard deviation.

Initialize the head's linear layer.

Parameters:

Name	Type	Description	Default
in_features	int	Number of input feature dimensions.	required
out_features	int	Number of predicted target dimensions.	required
rngs	Rngs	NNX RNG collection used to initialize parameters.	required
eps	float	A small positive constant added to the softplus-mapped scale to keep it strictly positive.	1e-06

Raises:

Type	Description
ValueError	If out_features or eps is invalid.

Source code in src/stagereg/jax/distributions.py

def __init__(
    self,
    in_features: int,
    out_features: int,
    *,
    rngs: nnx.Rngs,
    eps: float = 1e-6,
) -> None:
    """Initialize the head's linear layer.

    Args:
        in_features: Number of input feature dimensions.
        out_features: Number of predicted target dimensions.
        rngs: NNX RNG collection used to initialize parameters.
        eps: A small positive constant added to the ``softplus``-mapped
            scale to keep it strictly positive.

    Raises:
        ValueError: If ``out_features`` or ``eps`` is invalid.
    """
    if out_features <= 0:
        raise ValueError("out_features must be positive.")
    if not math.isfinite(eps) or eps <= 0.0:
        raise ValueError("eps must be positive and finite.")
    self.out_features = out_features
    self.eps = eps
    self.linear = nnx.Linear(in_features, 2 * out_features, rngs=rngs)

__call__

__call__(features: Array) -> Gaussian

Produce a Gaussian prediction from features.

Parameters:

Name	Type	Description	Default
features	Array	Model features, shaped (..., in_features).	required

Returns:

Type	Description
Gaussian	A Gaussian with loc/scale shaped
Gaussian	(..., out_features).

Source code in src/stagereg/jax/distributions.py

def __call__(self, features: jax.Array) -> Gaussian:
    """Produce a [`Gaussian`][stagereg.jax.Gaussian] prediction from ``features``.

    Args:
        features: Model features, shaped ``(..., in_features)``.

    Returns:
        A [`Gaussian`][stagereg.jax.Gaussian] with ``loc``/``scale`` shaped
        ``(..., out_features)``.
    """
    raw_loc, raw_scale = jnp.split(self.linear(features), 2, axis=-1)
    scale = jax.nn.softplus(raw_scale) + self.eps
    return Gaussian(loc=raw_loc, scale=scale)

GammaHead

GammaHead(
    in_features: int,
    out_features: int,
    *,
    rngs: Rngs,
    eps: float = 1e-06,
)

Bases: Module

An NNX head producing a shape/rate Gamma distribution.

Initialize the Gamma parameter projection.

Parameters:

Name	Type	Description	Default
in_features	int	Number of input feature dimensions.	required
out_features	int	Number of predicted target dimensions.	required
rngs	Rngs	NNX RNG collection used to initialize parameters.	required
eps	float	Positive finite offset added after softplus.	1e-06

Raises:

Type	Description
ValueError	If out_features or eps is invalid.

Source code in src/stagereg/jax/distributions.py

def __init__(
    self,
    in_features: int,
    out_features: int,
    *,
    rngs: nnx.Rngs,
    eps: float = 1e-6,
) -> None:
    """Initialize the Gamma parameter projection.

    Args:
        in_features: Number of input feature dimensions.
        out_features: Number of predicted target dimensions.
        rngs: NNX RNG collection used to initialize parameters.
        eps: Positive finite offset added after ``softplus``.

    Raises:
        ValueError: If ``out_features`` or ``eps`` is invalid.
    """
    if out_features <= 0:
        raise ValueError("out_features must be positive.")
    if not math.isfinite(eps) or eps <= 0.0:
        raise ValueError("eps must be positive and finite.")
    self.out_features = out_features
    self.eps = eps
    self.linear = nnx.Linear(in_features, 2 * out_features, rngs=rngs)

__call__

__call__(features: Array) -> Gamma

Produce a Gamma prediction from model features.

Parameters:

Name	Type	Description	Default
features	Array	Model features shaped (..., in_features).	required

Returns:

Type	Description
Gamma	A Gamma distribution with positive concentration and rate shaped
Gamma	(..., out_features).

Source code in src/stagereg/jax/distributions.py

def __call__(self, features: jax.Array) -> Gamma:
    """Produce a Gamma prediction from model features.

    Args:
        features: Model features shaped ``(..., in_features)``.

    Returns:
        A Gamma distribution with positive concentration and rate shaped
        ``(..., out_features)``.
    """
    raw_concentration, raw_rate = jnp.split(
        self.linear(features),
        2,
        axis=-1,
    )
    concentration = jax.nn.softplus(raw_concentration) + self.eps
    rate = jax.nn.softplus(raw_rate) + self.eps
    return Gamma(concentration=concentration, rate=rate)

stagereg.jax.evaluation

SupervisedLoss

Bases: Protocol

A callable computing a supervised loss for an NNX model.

Implementations compute a scalar loss from a model, a batch of inputs/targets/sample weights, a PRNG key (e.g. for stochastic layers such as dropout), and a flag indicating whether the call happens during training (as opposed to evaluation).

__call__

__call__(
    model: Module,
    inputs: PyTree,
    targets: Array,
    sample_weight: Array | None,
    key: Array,
    training: bool,
) -> jax.Array

Compute the scalar supervised loss for a batch.

Parameters:

Name	Type	Description	Default
model	Module	The NNX module to evaluate.	required
inputs	PyTree	The batch inputs.	required
targets	Array	The batch targets.	required
sample_weight	Array | None	Per-sample weights for the batch.	required
key	Array	A JAX PRNG key, e.g. for stochastic layers.	required
training	bool	Whether the loss is being computed during training (as opposed to evaluation).	required

Returns:

Type	Description
Array	The scalar loss value.

Source code in src/stagereg/jax/evaluation.py

def __call__(
    self,
    model: nnx.Module,
    inputs: PyTree,
    targets: jax.Array,
    sample_weight: jax.Array | None,
    key: jax.Array,
    training: bool,
    /,
) -> jax.Array:
    """Compute the scalar supervised loss for a batch.

    Args:
        model: The NNX module to evaluate.
        inputs: The batch inputs.
        targets: The batch targets.
        sample_weight: Per-sample weights for the batch.
        key: A JAX PRNG key, e.g. for stochastic layers.
        training: Whether the loss is being computed during training
            (as opposed to evaluation).

    Returns:
        The scalar loss value.
    """
    ...

make_evaluation_step

make_evaluation_step(
    loss: SupervisedLoss,
    *,
    metrics: Sequence[BatchMetricSpec] = (),
) -> Callable[..., Mapping[str, jax.Array]]

Create a JIT-compiled NNX supervised evaluation step.

Parameters:

Name	Type	Description	Default
loss	SupervisedLoss	A callable computing the supervised loss given a model, inputs, targets, sample weights, a PRNG key, and a training flag.	required
metrics	Sequence[BatchMetricSpec]	Registered JAX-native batch metrics.	()

Returns:

Type	Description
Callable[..., Mapping[str, Array]]	A JIT-compiled function ``evaluate_step(model, inputs, targets,
Callable[..., Mapping[str, Array]]	sample_weight, key)returning a mapping containing"loss"``
Callable[..., Mapping[str, Array]]	plus one scalar value per registered batch metric.

Source code in src/stagereg/jax/evaluation.py

def make_evaluation_step(
    loss: SupervisedLoss,
    *,
    metrics: Sequence[BatchMetricSpec] = (),
) -> Callable[..., Mapping[str, jax.Array]]:
    """Create a JIT-compiled NNX supervised evaluation step.

    Args:
        loss: A callable computing the supervised loss given a model,
            inputs, targets, sample weights, a PRNG key, and a
            ``training`` flag.
        metrics: Registered JAX-native batch metrics.

    Returns:
        A JIT-compiled function ``evaluate_step(model, inputs, targets,
        sample_weight, key)`` returning a mapping containing ``"loss"``
        plus one scalar value per registered batch metric.
    """

    @nnx.jit
    def evaluate_step(
        model: nnx.Module,
        inputs: PyTree,
        targets: jax.Array,
        sample_weight: jax.Array | None,
        key: jax.Array,
    ) -> Mapping[str, jax.Array]:
        loss_value = loss(model, inputs, targets, sample_weight, key, False)
        values: dict[str, jax.Array] = {"loss": loss_value}
        for spec in metrics:
            values[spec.name] = spec.metric(
                model,
                inputs,
                targets,
                sample_weight,
                key,
                False,
            )
        return values

    return evaluate_step

evaluate_loader

evaluate_loader(
    model: Module,
    loader: Iterable[Batch],
    *,
    key: Array,
    metrics: MetricSuite | None = None,
    loss: SupervisedLoss | None = None,
) -> tuple[dict[str, float], jax.Array]

Evaluate a loader and return reduced metrics and the advanced random key.

The evaluation step is built here from metrics and loss, so a caller that only wants to score a loader against metrics need not invent a loss to do it.

Parameters:

Name	Type	Description	Default
model	Module	The NNX module to evaluate.	required
loader	Iterable[Batch]	An iterable of batches to evaluate.	required
key	Array	The JAX PRNG key to use, advanced once per batch.	required
metrics	MetricSuite | None	Optional registered batch/epoch metrics for evaluation. When omitted, no metrics are collected.	None
loss	SupervisedLoss | None	Optional supervised loss to average across batches. When omitted, the returned mapping has no "loss" key.	None

Returns:

Type	Description
dict[str, float]	A tuple (metrics, key) where metrics contains the
Array	registered metrics plus "loss" if loss was given, and
tuple[dict[str, float], Array]	key is advanced past all consumed batches.

Source code in src/stagereg/jax/evaluation.py

def evaluate_loader(
    model: nnx.Module,
    loader: Iterable[Batch],
    *,
    key: jax.Array,
    metrics: MetricSuite | None = None,
    loss: SupervisedLoss | None = None,
) -> tuple[dict[str, float], jax.Array]:
    """Evaluate a loader and return reduced metrics and the advanced random key.

    The evaluation step is built here from ``metrics`` and ``loss``, so a
    caller that only wants to score a loader against metrics need not
    invent a loss to do it.

    Args:
        model: The NNX module to evaluate.
        loader: An iterable of batches to evaluate.
        key: The JAX PRNG key to use, advanced once per batch.
        metrics: Optional registered batch/epoch metrics for evaluation.
            When omitted, no metrics are collected.
        loss: Optional supervised loss to average across batches. When
            omitted, the returned mapping has no ``"loss"`` key.

    Returns:
        A tuple ``(metrics, key)`` where ``metrics`` contains the
        registered metrics plus ``"loss"`` if ``loss`` was given, and
        ``key`` is advanced past all consumed batches.
    """
    metric_suite = metrics if metrics is not None else MetricSuite()
    evaluation_step = (
        _make_metrics_step(metric_suite.batch)
        if loss is None
        else make_evaluation_step(loss, metrics=metric_suite.batch)
    )
    losses: list[float] | None = None if loss is None else []
    batch_metric_values = initialize_batch_metric_values(metric_suite.batch)
    epoch_metric_parts = [] if metric_suite.epoch else None

    for batch in loader:
        key, batch_key = split_key(key)
        step_output = evaluation_step(
            model,
            batch.inputs,
            batch.targets,
            batch.sample_weight,
            batch_key,
        )
        collect_step_metrics(
            step_output,
            metrics=metric_suite.batch,
            losses=losses,
            batch_metric_values=batch_metric_values,
            context="evaluation step",
        )
        maybe_collect_epoch_prediction_data(
            epoch_metric_parts,
            suite=metric_suite,
            model=model,
            batch=batch,
            batch_key=batch_key,
        )

    return (
        reduce_metric_suite(
            suite=metric_suite,
            losses=losses,
            batch_metric_values=batch_metric_values,
            epoch_metric_parts=epoch_metric_parts,
        ),
        key,
    )

stagereg.jax.losses

make_supervised_loss

make_supervised_loss(
    objective: DistributionLoss | PredictionLoss,
    *,
    reduction: Callable[[Array], Array] = jnp.mean,
) -> SupervisedLoss

Adapt a per-example objective to the JAX supervised-runner contract.

Parameters:

Name	Type	Description	Default
objective	DistributionLoss | PredictionLoss	Backend-neutral objective computing unreduced loss values from model predictions and a batch.	required
reduction	Callable[[Array], Array]	Callable reducing optionally weighted loss values to a scalar. Defaults to jnp.mean.	mean

Returns:

Type	Description
SupervisedLoss	A supervised loss that invokes the model, applies the objective and
SupervisedLoss	optional sample weights, then reduces to a scalar. Training calls use
SupervisedLoss	the live model in training mode. Evaluation calls use an isolated
SupervisedLoss	inference-mode clone so validation cannot mutate the live model.

Source code in src/stagereg/jax/losses.py

def make_supervised_loss(
    objective: DistributionLoss | PredictionLoss,
    *,
    reduction: Callable[[jax.Array], jax.Array] = jnp.mean,
) -> SupervisedLoss:
    """Adapt a per-example objective to the JAX supervised-runner contract.

    Args:
        objective: Backend-neutral objective computing unreduced loss values
            from model predictions and a batch.
        reduction: Callable reducing optionally weighted loss values to a
            scalar. Defaults to ``jnp.mean``.

    Returns:
        A supervised loss that invokes the model, applies the objective and
        optional sample weights, then reduces to a scalar. Training calls use
        the live model in training mode. Evaluation calls use an isolated
        inference-mode clone so validation cannot mutate the live model.
    """

    def supervised_loss(
        model: nnx.Module,
        inputs: PyTree,
        targets: jax.Array,
        sample_weight: jax.Array | None,
        key: jax.Array,
        training: bool,
    ) -> jax.Array:
        del key
        prediction_model = model
        if training:
            prediction_model.train()
        else:
            prediction_model = nnx.clone(model)
            prediction_model.eval()
        prediction = prediction_model(inputs)
        batch = Batch(inputs=inputs, targets=targets, sample_weight=sample_weight)
        per_example = objective.per_example(prediction, batch)
        if sample_weight is not None:
            sample_weight = align_sample_weight(sample_weight, per_example)
            per_example = per_example * sample_weight
        return reduction(per_example)

    return supervised_loss

stagereg.jax.metrics

BatchMetric

Bases: Protocol

A JAX-native metric computed per batch.

__call__

__call__(
    model: Module,
    inputs: PyTree,
    targets: Array,
    sample_weight: Array | None,
    key: Array,
    training: bool,
) -> jax.Array

Return a scalar batch metric value.

Source code in src/stagereg/jax/metrics.py

def __call__(
    self,
    model: nnx.Module,
    inputs: PyTree,
    targets: jax.Array,
    sample_weight: jax.Array | None,
    key: jax.Array,
    training: bool,
    /,
) -> jax.Array:
    """Return a scalar batch metric value."""
    ...

BatchMetricSpec dataclass

BatchMetricSpec(
    name: str,
    metric: BatchMetric,
    reduce: Callable[
        [Sequence[float]], float
    ] = _metric_mean,
)

A named batch metric with a host-side reduction policy.

Attributes:

Name	Type	Description
name	str	Bare metric name the reduced value is reported under. Must not be empty.
metric	BatchMetric	The BatchMetric evaluated on every batch.
reduce	Callable[[Sequence[float]], float]	Host-side reduction of the per-batch values to one epoch value. Defaults to the arithmetic mean.

PredictionRequirements dataclass

PredictionRequirements(
    fields: MetricRequirements,
    predictive_sample_count: int | None = None,
    evaluation_grid: EvaluationGrid | None = None,
)

Combined metric requirements plus explicit materialization configuration.

Attributes:

Name	Type	Description
fields	MetricRequirements	Union of the prediction fields the registered epoch metrics read.
predictive_sample_count	int | None	Number of predictive samples per scoring unit. Required, and positive, when fields asks for predictive samples.
evaluation_grid	EvaluationGrid | None	Grid the CRPS metrics integrate over. Required when fields asks for an evaluation grid.

CoordinateExtractor

Bases: Protocol

Extract a numeric scoring-unit coordinate from a domain batch.

__call__

__call__(batch: Batch) -> Any

Return coordinates corresponding to the batch targets.

Source code in src/stagereg/jax/metrics.py

def __call__(self, batch: Batch, /) -> Any:
    """Return coordinates corresponding to the batch targets."""
    ...

ReferenceSamplesExtractor

Bases: Protocol

Materialize reference samples at the JAX/domain boundary.

__call__

__call__(batch: Batch, key: Array) -> Any

Return samples shaped (n_scoring_units, n_samples).

Source code in src/stagereg/jax/metrics.py

def __call__(self, batch: Batch, key: jax.Array, /) -> Any:
    """Return samples shaped ``(n_scoring_units, n_samples)``."""
    ...

Predictor

Bases: Protocol

Materialize typed host predictions from a JAX model and batch.

__call__

__call__(
    model: Module,
    batch: Batch,
    requirements: PredictionRequirements,
    key: Array,
) -> EpochPredictionData

Produce host-resident prediction data for batch.

Source code in src/stagereg/jax/metrics.py

def __call__(
    self,
    model: nnx.Module,
    batch: Batch,
    requirements: PredictionRequirements,
    key: jax.Array,
    /,
) -> EpochPredictionData:
    """Produce host-resident prediction data for ``batch``."""
    ...

GaussianPredictor dataclass

GaussianPredictor(
    coordinate_extractor: CoordinateExtractor | None = None,
    reference_samples_extractor: ReferenceSamplesExtractor
    | None = None,
)

Materialize scalar epoch data from a model returning Gaussian.

Attributes:

Name	Type	Description
coordinate_extractor	CoordinateExtractor | None	Optional explicit adapter for coordinate metadata.
reference_samples_extractor	ReferenceSamplesExtractor | None	Optional explicit adapter for conditional CRPS reference samples. It must return (n_scoring_units, n_samples).

__call__

__call__(
    model: Module,
    batch: Batch,
    requirements: PredictionRequirements,
    key: Array,
) -> EpochPredictionData

Materialize only fields requested by registered epoch metrics.

Parameters:

Name	Type	Description	Default
model	Module	Model mapping batch inputs to a scalar-event Gaussian.	required
batch	Batch	Domain batch containing targets.	required
requirements	PredictionRequirements	Combined fields and explicit sample/grid configuration.	required
key	Array	Dedicated metric PRNG key.	required

Returns:

Type	Description
EpochPredictionData	Validated host-resident prediction data.

Raises:

Type	Description
TypeError	If the model does not return a Gaussian.
ValueError	If targets, event shape, extractors, or shapes are invalid.

Source code in src/stagereg/jax/metrics.py

def __call__(
    self,
    model: nnx.Module,
    batch: Batch,
    requirements: PredictionRequirements,
    key: jax.Array,
    /,
) -> EpochPredictionData:
    """Materialize only fields requested by registered epoch metrics.

    Args:
        model: Model mapping batch inputs to a scalar-event Gaussian.
        batch: Domain batch containing targets.
        requirements: Combined fields and explicit sample/grid configuration.
        key: Dedicated metric PRNG key.

    Returns:
        Validated host-resident prediction data.

    Raises:
        TypeError: If the model does not return a Gaussian.
        ValueError: If targets, event shape, extractors, or shapes are invalid.
    """
    targets = _require_targets(batch)
    prediction = _require_scalar_gaussian(model(batch.inputs))
    return _materialize_epoch_data(
        prediction,
        targets=targets,
        batch=batch,
        requirements=requirements,
        key=key,
        interval=_gaussian_interval,
        coordinate_extractor=self.coordinate_extractor,
        reference_samples_extractor=self.reference_samples_extractor,
    )

PosteriorPredictivePredictor dataclass

PosteriorPredictivePredictor(
    coordinate_extractor: CoordinateExtractor | None = None,
    reference_samples_extractor: ReferenceSamplesExtractor
    | None = None,
)

Materialize scalar epoch data from a model returning PosteriorPredictive.

Scores the exact mixture rather than its moment-matched Gaussian: the log-likelihood is the mixture log-density, predictive samples are drawn from the mixture, and prediction intervals are central mixture quantiles.

Attributes:

Name	Type	Description
coordinate_extractor	CoordinateExtractor | None	Optional explicit adapter for coordinate metadata.
reference_samples_extractor	ReferenceSamplesExtractor | None	Optional explicit adapter for conditional CRPS reference samples. It must return (n_scoring_units, n_samples).

__call__

__call__(
    model: Module,
    batch: Batch,
    requirements: PredictionRequirements,
    key: Array,
) -> EpochPredictionData

Materialize only fields requested by registered epoch metrics.

Parameters:

Name	Type	Description	Default
model	Module	Model mapping batch inputs to a PosteriorPredictive.	required
batch	Batch	Domain batch containing targets.	required
requirements	PredictionRequirements	Combined fields and explicit sample/grid configuration.	required
key	Array	Dedicated metric PRNG key.	required

Returns:

Type	Description
EpochPredictionData	Validated host-resident prediction data.

Raises:

Type	Description
TypeError	If the model does not return a posterior predictive.
ValueError	If targets, extractors, or shapes are invalid.

Source code in src/stagereg/jax/metrics.py

def __call__(
    self,
    model: nnx.Module,
    batch: Batch,
    requirements: PredictionRequirements,
    key: jax.Array,
    /,
) -> EpochPredictionData:
    """Materialize only fields requested by registered epoch metrics.

    Args:
        model: Model mapping batch inputs to a
            [`PosteriorPredictive`][stagereg.jax.PosteriorPredictive].
        batch: Domain batch containing targets.
        requirements: Combined fields and explicit sample/grid configuration.
        key: Dedicated metric PRNG key.

    Returns:
        Validated host-resident prediction data.

    Raises:
        TypeError: If the model does not return a posterior predictive.
        ValueError: If targets, extractors, or shapes are invalid.
    """
    targets = _require_targets(batch)
    prediction = model(batch.inputs)
    if not isinstance(prediction, PosteriorPredictive):
        raise TypeError(
            "PosteriorPredictivePredictor requires a model returning "
            "PosteriorPredictive."
        )
    return _materialize_epoch_data(
        prediction,
        targets=targets,
        batch=batch,
        requirements=requirements,
        key=key,
        interval=_mixture_interval,
        coordinate_extractor=self.coordinate_extractor,
        reference_samples_extractor=self.reference_samples_extractor,
    )

MetricSuite dataclass

MetricSuite(
    batch: tuple[BatchMetricSpec, ...] = (),
    epoch: tuple[EpochMetric, ...] = (),
    predictor: Predictor | None = None,
    predictive_sample_count: int | None = None,
    evaluation_grid: EvaluationGrid | None = None,
)

Complete metric registration and materialization configuration.

Attributes:

Name	Type	Description
batch	tuple[BatchMetricSpec, ...]	Batch metrics evaluated on every batch.
epoch	tuple[EpochMetric, ...]	EpochMetrics evaluated once per epoch over the merged predictions.
predictor	Predictor | None	The Predictor that turns model outputs into host arrays. Required when epoch is non-empty.
predictive_sample_count	int | None	Number of predictive samples per scoring unit. Required by sample-based epoch metrics.
evaluation_grid	EvaluationGrid | None	Grid the CRPS metrics integrate over. Required by CRPS epoch metrics.

metric_requirements property

metric_requirements: MetricRequirements

Return the union of registered epoch metric requirements.

prediction_requirements property

prediction_requirements: PredictionRequirements

Return combined fields plus explicit suite configuration.

merge_epoch_prediction_data

merge_epoch_prediction_data(
    parts: Sequence[EpochPredictionData],
) -> EpochPredictionData

Merge batches according to each typed field's scoring axes.

Parameters:

Name	Type	Description	Default
parts	Sequence[EpochPredictionData]	Per-batch host-materialized prediction data.	required

Returns:

Type	Description
EpochPredictionData	One validated epoch prediction object.

Raises:

Type	Description
ValueError	If parts are empty or optional fields/configuration conflict.

Source code in src/stagereg/jax/metrics.py

def merge_epoch_prediction_data(
    parts: Sequence[EpochPredictionData],
) -> EpochPredictionData:
    """Merge batches according to each typed field's scoring axes.

    Args:
        parts: Per-batch host-materialized prediction data.

    Returns:
        One validated epoch prediction object.

    Raises:
        ValueError: If parts are empty or optional fields/configuration conflict.
    """
    if not parts:
        raise ValueError("at least one epoch prediction batch is required.")
    interval_levels = {interval.level for interval in parts[0].intervals}
    if any(
        {interval.level for interval in part.intervals} != interval_levels
        for part in parts
    ):
        raise ValueError("prediction interval levels must match across batches.")
    intervals = tuple(
        PredictionInterval(
            level=level,
            lower=np.concatenate([part.interval(level).lower for part in parts]),
            upper=np.concatenate([part.interval(level).upper for part in parts]),
        )
        for level in sorted(interval_levels)
    )
    return EpochPredictionData(
        targets=np.concatenate([part.targets for part in parts]),
        mean=np.concatenate([part.mean for part in parts]),
        variance=_concat_optional_vectors(parts, "variance"),
        predictive_samples=_concat_optional_samples(parts, "predictive_samples"),
        reference_samples=_concat_optional_samples(parts, "reference_samples"),
        intervals=intervals,
        coordinate=_concat_optional_vectors(parts, "coordinate"),
        evaluation_grid=_merge_grid(parts),
        log_likelihood=_concat_optional_vectors(parts, "log_likelihood"),
    )

stagereg.jax.posterior

PosteriorProblem dataclass

PosteriorProblem(
    initial_parameters: PyTree,
    mean_function: Callable[[PyTree, PyTree], Array],
    log_likelihood: Callable[[PyTree, Batch], Array],
    prior: Prior,
    train_loader: LoaderFactory,
    dataset_size: int,
)

Everything a posterior stage hands its inference method.

Attributes:

Name	Type	Description
initial_parameters	PyTree	The warm start: a copy of the trained mean network's parameters, in the posterior network's parameter tree.
mean_function	Callable[[PyTree, PyTree], Array]	Maps (parameters, inputs) to the posterior network's mean predictions at inputs.
log_likelihood	Callable[[PyTree, Batch], Array]	Maps (parameters, batch) to the scalar Gaussian log-likelihood of the batch targets around the mean predictions, under the fixed aleatoric variance, scaled by N / B so that a minibatch estimates the full-data log-likelihood. Differentiable in parameters; the aleatoric variance carries no gradient.
prior	Prior	The prior over the parameter tree.
train_loader	LoaderFactory	Factory producing the training batches of each epoch.
dataset_size	int	The number of training examples N.

Prior

Bases: Protocol

A log-density over the parameter tree of the posterior network.

log_prob

log_prob(parameters: PyTree) -> jax.Array

Return the scalar log-density of parameters.

Parameters:

Name	Type	Description	Default
parameters	PyTree	A tree with the posterior network's parameter tree.	required

Returns:

Type	Description
Array	The scalar prior log-density.

Source code in src/stagereg/jax/posterior.py

def log_prob(self, parameters: PyTree) -> jax.Array:
    """Return the scalar log-density of ``parameters``.

    Args:
        parameters: A tree with the posterior network's parameter tree.

    Returns:
        The scalar prior log-density.
    """
    ...

IsotropicGaussianPrior dataclass

IsotropicGaussianPrior(precision: float = 1.0)

An independent zero-mean Gaussian on every parameter.

The default precision of 1 is VeBNN's κ = 1.

Attributes:

Name	Type	Description
precision	float	The inverse variance of every parameter, positive and finite. Defaults to 1.0.

log_prob

log_prob(parameters: PyTree) -> jax.Array

Return the normalized log-density of a parameter tree.

Parameters:

Name	Type	Description	Default
parameters	PyTree	A tree of parameter arrays.	required

Returns:

Type	Description
Array	The scalar sum, over every parameter, of its Gaussian log-density.

Source code in src/stagereg/jax/posterior.py

def log_prob(self, parameters: PyTree) -> jax.Array:
    """Return the normalized log-density of a parameter tree.

    Args:
        parameters: A tree of parameter arrays.

    Returns:
        The scalar sum, over every parameter, of its Gaussian log-density.
    """
    leaves = jax.tree.leaves(parameters)
    count = sum(jnp.size(leaf) for leaf in leaves)
    squared_norm = sum(jnp.sum(jnp.square(leaf)) for leaf in leaves)
    normalizer = 0.5 * count * math.log(self.precision / (2.0 * math.pi))
    return normalizer - 0.5 * self.precision * jnp.asarray(squared_norm)

InferenceMethod

Bases: Protocol

How a posterior stage turns a posterior problem into a posterior.

The stage owns the epoch loop: it calls init once, then update once per training batch, and reads the current posterior to validate each epoch that ends with has_posterior true. Every method honors the problem's prior.

supports_early_stopping property

supports_early_stopping: bool

Whether stopping on a validation improvement is meaningful.

False for SG-MCMC, whose chain an early stopper would truncate.

has_posterior property

has_posterior: bool

Whether posterior can be read.

False before init and, for SG-MCMC, until the first sample is retained. The posterior stage skips validation for an epoch that ends without a posterior, so a method that supports early stopping must have one after every epoch.

init

init(problem: PosteriorProblem) -> None

Start inference from problem, discarding any previous state.

Parameters:

Name	Type	Description	Default
problem	PosteriorProblem	The posterior problem.	required

Source code in src/stagereg/jax/posterior.py

def init(self, problem: PosteriorProblem) -> None:
    """Start inference from ``problem``, discarding any previous state.

    Args:
        problem: The posterior problem.
    """
    ...

update

update(batch: Batch, key: Array) -> Mapping[str, Any]

Advance inference by one training batch.

Parameters:

Name	Type	Description	Default
batch	Batch	The training batch.	required
key	Array	A fresh PRNG key for this batch.	required

Returns:

Type	Description
Mapping[str, Any]	A mapping holding the scalar "loss" the method minimizes.

Source code in src/stagereg/jax/posterior.py

def update(self, batch: Batch, key: jax.Array) -> Mapping[str, Any]:
    """Advance inference by one training batch.

    Args:
        batch: The training batch.
        key: A fresh PRNG key for this batch.

    Returns:
        A mapping holding the scalar ``"loss"`` the method minimizes.
    """
    ...

posterior

posterior() -> Posterior

Return the current posterior, independent of later updates.

Returns:

Type	Description
Posterior	The posterior inferred so far.

Source code in src/stagereg/jax/posterior.py

def posterior(self) -> Posterior:
    """Return the current posterior, independent of later updates.

    Returns:
        The posterior inferred so far.
    """
    ...

state

state() -> PyTree

Return the method's full state, independent of later updates.

This is what a best checkpoint stores, so that inference can resume from it, e.g. variational parameters plus optimizer state.

Returns:

Type	Description
PyTree	A tree of arrays.

Source code in src/stagereg/jax/posterior.py

def state(self) -> PyTree:
    """Return the method's full state, independent of later updates.

    This is what a best checkpoint stores, so that inference can resume
    from it, e.g. variational parameters plus optimizer state.

    Returns:
        A tree of arrays.
    """
    ...

load_state

load_state(state: PyTree) -> None

Resume inference from a state returned by state().

Parameters:

Name	Type	Description	Default
state	PyTree	A state previously returned by state.	required

Source code in src/stagereg/jax/posterior.py

def load_state(self, state: PyTree) -> None:
    """Resume inference from a state returned by ``state()``.

    Args:
        state: A state previously returned by
            [`state`][stagereg.jax.InferenceMethod.state].
    """
    ...

posterior_state

posterior_state() -> PyTree

Return the current posterior alone, independent of later updates.

This is what a finalized checkpoint stores: only what load_posterior needs to rebuild the posterior, e.g. the variational parameters of a variational method or the retained samples of an SG-MCMC method, but no optimizer state.

Returns:

Type	Description
PyTree	A tree of arrays.

Source code in src/stagereg/jax/posterior.py

def posterior_state(self) -> PyTree:
    """Return the current posterior alone, independent of later updates.

    This is what a finalized checkpoint stores: only what
    [`load_posterior`][stagereg.jax.InferenceMethod.load_posterior] needs
    to rebuild the posterior, e.g. the variational parameters of a
    variational method or the retained samples of an SG-MCMC method, but
    no optimizer state.

    Returns:
        A tree of arrays.
    """
    ...

load_posterior

load_posterior(state: PyTree) -> None

Rebuild a finalized posterior after init.

Afterwards posterior returns the posterior that state describes. The method need not be able to resume inference from it.

Parameters:

Name	Type	Description	Default
state	PyTree	A state previously returned by posterior_state.	required

Source code in src/stagereg/jax/posterior.py

def load_posterior(self, state: PyTree) -> None:
    """Rebuild a finalized posterior after
    [`init`][stagereg.jax.InferenceMethod.init].

    Afterwards [`posterior`][stagereg.jax.InferenceMethod.posterior]
    returns the posterior that ``state`` describes. The method need not be
    able to resume inference from it.

    Args:
        state: A state previously returned by
            [`posterior_state`][stagereg.jax.InferenceMethod.posterior_state].
    """
    ...

Posterior

Bases: Protocol

An approximate posterior over mean functions, consumed through draws.

A draw is a whole mean function: under the same key, draw s is the same function at every input and in every batch. A finite posterior (retained samples, ensemble members) has a fixed set of draws and always uses all of them; an unlimited one (a variational distribution, MC-Dropout) takes the number of draws from its caller.

num_draws property

num_draws: int | None

The number of draws of a finite posterior, or None if unlimited.

sample_means

sample_means(
    inputs: PyTree,
    key: Array,
    num_samples: int | None = None,
) -> jax.Array

Draw mean predictions at inputs.

Parameters:

Name	Type	Description	Default
inputs	PyTree	Model inputs with a leading batch dimension.	required
key	Array	The PRNG key selecting the draws. A finite posterior may ignore it.	required
num_samples	int | None	The number of draws S. Must be None for a finite posterior and given for an unlimited one.	None

Returns:

Type	Description
Array	The draws' mean predictions, shaped [S, *batch].

Raises:

Type	Description
ValueError	If num_samples is given to a finite posterior, or omitted for an unlimited one.

Source code in src/stagereg/jax/posterior.py

def sample_means(
    self,
    inputs: PyTree,
    key: jax.Array,
    num_samples: int | None = None,
) -> jax.Array:
    """Draw mean predictions at ``inputs``.

    Args:
        inputs: Model inputs with a leading batch dimension.
        key: The PRNG key selecting the draws. A finite posterior may
            ignore it.
        num_samples: The number of draws ``S``. Must be ``None`` for a
            finite posterior and given for an unlimited one.

    Returns:
        The draws' mean predictions, shaped ``[S, *batch]``.

    Raises:
        ValueError: If ``num_samples`` is given to a finite posterior, or
            omitted for an unlimited one.
    """
    ...

stagereg.jax.posterior_stage

PosteriorStageOptions dataclass

PosteriorStageOptions(
    epochs: int,
    num_draws: int = 32,
    validation_loader: LoaderFactory | None = None,
    validation_metrics: MetricSuite = _default_validation_metrics(),
    early_stopper: EarlyStopper | None = None,
    event_sinks: Sequence[EventSink] = (),
    checkpoint_store: CheckpointStore | None = None,
    checkpoint_key: str | None = None,
)

Epoch-loop and prediction options of a posterior stage.

Attributes:

Name	Type	Description
epochs	int	Maximum number of training epochs.
num_draws	int	The number of draws S taken from an unlimited posterior (e.g. a variational distribution) to form the posterior predictive. A finite posterior always uses all its draws and ignores it. Defaults to 32.
validation_loader	LoaderFactory | None	Factory producing the "validation" batches the posterior predictive is scored on after each epoch. Defaults to None, which disables validation.
validation_metrics	MetricSuite	The epoch metrics the posterior predictive is scored with, reported under their usual names (nll, crps, ...). Its predictor receives a model mapping inputs to a PosteriorPredictive, so it is normally a PosteriorPredictivePredictor. Defaults to the exact NLL (nll) and the CRPS (crps) of 100 predictive samples per target, scored with SampleContinuousRankedProbabilityScore, which needs no evaluation grid.
early_stopper	EarlyStopper | None	Optional early-stopping policy. Refused for an inference method that does not support early stopping.
event_sinks	Sequence[EventSink]	Event consumers notified of posterior/... events.
checkpoint_store	CheckpointStore | None	Optional store for the best checkpoint, which holds the inference method's full state, and for the finalized checkpoint that replaces it at the end of training and holds the posterior alone.
checkpoint_key	str | None	The key of the best and finalized checkpoints. Defaults to None, which resolves to posterior/best.

PosteriorStage dataclass

PosteriorStage(
    inference_method: InferenceMethod,
    train_loader: LoaderFactory,
    dataset_size: int,
    options: PosteriorStageOptions,
    prior: Prior = IsotropicGaussianPrior(),
    model: Module | None = None,
    mean_model_name: str = "mean_model",
    variance_model_name: str = "variance_model",
    model_name: str = "posterior",
)

Optional Step 3 stage inferring a posterior over mean functions.

It follows the variance stage. The posterior network has the mean network's parameter tree and is warm-started from a copy of the trained mean weights; the mean and variance models stay frozen, and the aleatoric variance is computed on the fly as the mean of the variance model's prediction. The stage hands its inference method a PosteriorProblem, drives the epoch loop with the method's update, and validates every epoch on the current posterior predictive, skipping epochs that end before the method has_posterior.

Attributes:

Name	Type	Description
inference_method	InferenceMethod	The InferenceMethod deciding how the posterior is inferred.
train_loader	LoaderFactory	Factory producing the original regression batches.
dataset_size	int	The number of training examples N, which scales each batch's log-likelihood by N / B.
options	PosteriorStageOptions	Epoch-loop and prediction options.
prior	Prior	The prior over the posterior network's parameters. Defaults to IsotropicGaussianPrior with precision 1.
model	Module | None	The posterior network. A copy of it is warm-started, so its own parameters never change. Its parameter tree must equal the mean network's, but its module may differ. Defaults to None, which uses a copy of the mean network.
mean_model_name	str	Registry name of the trained mean model.
variance_model_name	str	Registry name of the trained variance model.
model_name	str	Registry name the posterior is registered under.

prepare

prepare(state: TrainingState) -> None

Warm-start the posterior network and initialize the inference method.

Every check runs before state changes. Afterwards the mean and variance models are frozen components.

Parameters:

Name	Type	Description	Default
state	TrainingState	Shared state whose mean and variance stages are ready.	required

Raises:

Type	Description
ValueError	If the lifecycle state is not VARIANCE_READY, if the mean or variance model is not registered with its role, if model_name is taken by another component, if the posterior network's parameter tree differs from the mean network's, if dataset_size is not positive, if an early stopper is configured for a method that does not support early stopping, or if the epoch-loop configuration is refused.
TypeError	If the mean or variance model is not an NNX module, or state.rng_state is not a JAX random key.

Source code in src/stagereg/jax/posterior_stage.py

def prepare(self, state: TrainingState) -> None:
    """Warm-start the posterior network and initialize the inference method.

    Every check runs before ``state`` changes. Afterwards the mean and
    variance models are frozen components.

    Args:
        state: Shared state whose mean and variance stages are ready.

    Raises:
        ValueError: If the lifecycle state is not ``VARIANCE_READY``, if
            the mean or variance model is not registered with its role, if
            ``model_name`` is taken by another component, if the posterior
            network's parameter tree differs from the mean network's, if
            ``dataset_size`` is not positive, if an early stopper is
            configured for a method that does not support early stopping,
            or if the epoch-loop configuration is refused.
        TypeError: If the mean or variance model is not an NNX module, or
            ``state.rng_state`` is not a JAX random key.
    """
    if state.lifecycle_state is not StageState.VARIANCE_READY:
        raise ValueError("posterior stage requires VARIANCE_READY lifecycle state.")
    mean_model = _registered_module(
        state, self.mean_model_name, ParameterRole.MEAN, kind="mean"
    )
    variance_model = _registered_module(
        state, self.variance_model_name, ParameterRole.VARIANCE, kind="variance"
    )
    if self.model_name in state.model_components:
        raise ValueError(
            f"model component {self.model_name!r} is already registered."
        )
    if self.dataset_size <= 0:
        raise ValueError("dataset_size must be positive.")
    options = self.options
    if (
        options.early_stopper is not None
        and not self.inference_method.supports_early_stopping
    ):
        raise ValueError(
            "the inference method does not support early stopping; "
            "remove the early stopper."
        )
    validation = (
        None
        if options.validation_loader is None
        else _PosteriorPredictiveValidation(
            inference_method=self.inference_method,
            aleatoric_variance=_aleatoric_variance(variance_model),
            loader=options.validation_loader,
            metrics=options.validation_metrics,
            num_draws=options.num_draws,
        )
    )
    check_epoch_loop_arguments(
        state=state,
        epochs=options.epochs,
        stage=self.name,
        validation=validation,
        early_stopper=options.early_stopper,
    )
    self.inference_method.init(self._problem(mean_model, variance_model))
    state.frozen_components = state.frozen_components | {
        self.mean_model_name,
        self.variance_model_name,
    }
    state.active_stage = self.name
    self._validation = validation
    self._prepared = True
    self._posterior = None

train

train(state: TrainingState) -> StageResult

Infer the posterior, register it and transition to POSTERIOR_READY.

With an early stopper and a checkpoint store, every improvement saves a best checkpoint whose parameters hold the inference method's full state. After the last epoch the method resumes from the best checkpoint, if one was saved, and state returns to the best epoch. The method's posterior_state is then stored as state.posterior_state, and, with a checkpoint store, saved as the finalized checkpoint under the checkpoint key: lifecycle state POSTERIOR_READY, metadata {"stage": "posterior", "stage_complete": True}, and no parameters, so neither the mean and variance weights nor the method's optimizer state are duplicated.

Parameters:

Name	Type	Description	Default
state	TrainingState	State the stage has been prepared on.	required

Returns:

Type	Description
StageResult	The epoch loop's result: the training metrics of the best epoch
StageResult	when a best checkpoint was restored, otherwise of the last epoch.

Raises:

Type	Description
ValueError	If the stage was not prepared or training produced a non-finite final loss.

Source code in src/stagereg/jax/posterior_stage.py

def train(self, state: TrainingState) -> StageResult:
    """Infer the posterior, register it and transition to ``POSTERIOR_READY``.

    With an early stopper and a checkpoint store, every improvement saves a
    best checkpoint whose ``parameters`` hold the inference method's full
    [`state`][stagereg.jax.InferenceMethod.state]. After the last epoch the
    method resumes from the best checkpoint, if one was saved, and
    ``state`` returns to the best epoch. The method's
    [`posterior_state`][stagereg.jax.InferenceMethod.posterior_state] is
    then stored as ``state.posterior_state``, and, with a checkpoint store,
    saved as the finalized checkpoint under the checkpoint key: lifecycle
    state ``POSTERIOR_READY``, metadata ``{"stage": "posterior",
    "stage_complete": True}``, and no parameters, so neither the mean and
    variance weights nor the method's optimizer state are duplicated.

    Args:
        state: State the stage has been prepared on.

    Returns:
        The epoch loop's result: the training metrics of the best epoch
        when a best checkpoint was restored, otherwise of the last epoch.

    Raises:
        ValueError: If the stage was not prepared or training produced a
            non-finite final loss.
    """
    if state.lifecycle_state is not StageState.VARIANCE_READY or not self._prepared:
        raise ValueError("posterior stage must be prepared before training.")
    method = self.inference_method
    options = self.options
    result = run_epoch_loop(
        step=method.update,
        snapshot_state=lambda: StateSnapshot(parameters=method.state()),
        train_loader=self.train_loader,
        state=state,
        epochs=options.epochs,
        stage=self.name,
        checkpoint_key=options.checkpoint_key,
        validation=self._validation,
        early_stopper=options.early_stopper,
        event_sinks=options.event_sinks,
        checkpoint_store=options.checkpoint_store,
    )
    if result.loss is None or not math.isfinite(result.loss):
        raise ValueError(f"{self.name} stage produced a non-finite final loss.")
    result = self._finalize(state, result)
    self._prepared = False
    return result

restore

restore(
    state: TrainingState, checkpoint: Checkpoint
) -> None

Restore the posterior stage's finalized checkpoint into state.

Call it after the mean stage's restore and the variance stage's restore. The inference method is initialized on the posterior problem of the registered mean and variance models, then rebuilds the saved posterior with load_posterior; the posterior is registered under model_name. Every model component and optimizer registered before, such as the mean and variance models, stays registered. Afterwards the state passes validate.

Parameters:

Name	Type	Description	Default
state	TrainingState	Training state the mean and variance stages have been restored into.	required
checkpoint	Checkpoint	The posterior stage's finalized checkpoint.	required

Raises:

Type	Description
ValueError	If checkpoint is not finalized by this stage, that is its lifecycle state is not POSTERIOR_READY, its metadata lacks "stage": "posterior" and "stage_complete": True, or it holds no posterior state; if it was saved with another model_name, mean_model_name or variance_model_name; if the mean or variance model is not registered with its role; if dataset_size is not positive; or if the posterior network's parameter tree differs from the mean network's; or if the inference method's load_posterior refuses the saved posterior state. All are checked before state, the inference method or model changes.
TypeError	If the mean or variance model is not an NNX module, or the checkpoint holds no JAX random key.

Source code in src/stagereg/jax/posterior_stage.py

def restore(self, state: TrainingState, checkpoint: Checkpoint) -> None:
    """Restore the posterior stage's finalized checkpoint into ``state``.

    Call it after the mean stage's
    [`restore`][stagereg.jax.MeanStage.restore] and the variance stage's
    [`restore`][stagereg.jax.GammaVarianceStage.restore]. The inference
    method is initialized on the posterior problem of the registered mean
    and variance models, then rebuilds the saved posterior with
    [`load_posterior`][stagereg.jax.InferenceMethod.load_posterior]; the
    posterior is registered under ``model_name``. Every model component
    and optimizer registered before, such as the mean and variance models,
    stays registered. Afterwards the state passes
    [`validate`][stagereg.jax.PosteriorStage.validate].

    Args:
        state: Training state the mean and variance stages have been
            restored into.
        checkpoint: The posterior stage's finalized checkpoint.

    Raises:
        ValueError: If ``checkpoint`` is not finalized by this stage, that
            is its lifecycle state is not ``POSTERIOR_READY``, its metadata
            lacks ``"stage": "posterior"`` and ``"stage_complete": True``,
            or it holds no posterior state; if it was saved with another
            ``model_name``, ``mean_model_name`` or ``variance_model_name``;
            if the mean or variance model is not registered with its role;
            if ``dataset_size`` is not positive; or if the posterior
            network's parameter tree differs from the mean network's; or
            if the inference method's ``load_posterior`` refuses the saved
            posterior state. All are checked before ``state``, the inference
            method or ``model`` changes.
        TypeError: If the mean or variance model is not an NNX module, or
            the checkpoint holds no JAX random key.
    """
    posterior_state, problem = self._restorable(state, checkpoint)
    # A dry run on a copy, so a refused posterior state leaves the method as is.
    trial = copy.deepcopy(self.inference_method)
    trial.init(problem)
    trial.load_posterior(posterior_state)

    self.inference_method.init(problem)
    self.inference_method.load_posterior(posterior_state)
    restore_keeping_components(state, checkpoint, drop={self.model_name})
    self._register_posterior(state, posterior_state)
    self._validation = None
    self._prepared = False

validate

validate(state: TrainingState) -> ValidationResult

Validate posterior-stage lifecycle, ownership and freezing.

Parameters:

Name	Type	Description	Default
state	TrainingState	Shared staged training state.	required

Returns:

Type	Description
ValidationResult	A validation result describing whether Step 3 is ready.

Source code in src/stagereg/jax/posterior_stage.py

def validate(self, state: TrainingState) -> ValidationResult:
    """Validate posterior-stage lifecycle, ownership and freezing.

    Args:
        state: Shared staged training state.

    Returns:
        A validation result describing whether Step 3 is ready.
    """
    passed = (
        state.lifecycle_state is StageState.POSTERIOR_READY
        and self._posterior is not None
        and state.model_components.get(self.model_name) is self._posterior
        and state.parameter_roles.get(self.model_name) is ParameterRole.POSTERIOR
        and state.parameter_roles.get(self.mean_model_name) is ParameterRole.MEAN
        and state.parameter_roles.get(self.variance_model_name)
        is ParameterRole.VARIANCE
        and self.mean_model_name in state.model_components
        and self.variance_model_name in state.model_components
        and {self.mean_model_name, self.variance_model_name}
        <= state.frozen_components
    )
    return ValidationResult(
        passed=passed,
        message=None if passed else "posterior stage invariants are not satisfied.",
    )

select_checkpoint

select_checkpoint(state: TrainingState) -> CheckpointRef

Return a reference to the posterior stage's best checkpoint.

Parameters:

Name	Type	Description	Default
state	TrainingState	Shared staged training state.	required

Returns:

Type	Description
CheckpointRef	Reference to the configured checkpoint_key, or
CheckpointRef	posterior/best when none was configured.

Raises:

Type	Description
ValueError	If no checkpoint exists under that key.

Source code in src/stagereg/jax/posterior_stage.py

def select_checkpoint(self, state: TrainingState) -> CheckpointRef:
    """Return a reference to the posterior stage's best checkpoint.

    Args:
        state: Shared staged training state.

    Returns:
        Reference to the configured ``checkpoint_key``, or
        ``posterior/best`` when none was configured.

    Raises:
        ValueError: If no checkpoint exists under that key.
    """
    del state
    return select_checkpoint(
        self.options.checkpoint_store,
        key=resolve_checkpoint_key(self.options.checkpoint_key, self.name),
        stage=self.name,
    )

stagereg.jax.psgld

PreconditionedSGLD dataclass

PreconditionedSGLD(
    step_size: float,
    burn_in: int = 0,
    thinning: int = 1,
    decay: float = 0.99,
    stability: float = 1e-05,
)

RMSprop-preconditioned stochastic-gradient Langevin dynamics.

Each update takes one step on the batch's estimate U of the negative log posterior, the problem's log-likelihood (already scaled to the full data set) plus its prior. With gradient g = ∇U, the preconditioner tracks V ← decay · V + (1 − decay) · g² and G = 1 / (stability + √V), and the parameters move by −(step_size / 2) · G · g plus Gaussian noise of variance step_size · G. The preconditioner's own drift correction is neglected, as is usual.

The step is compiled once per init; the log-likelihood sees each batch's inputs, targets and sample weights but not its metadata.

Burn-in and thinning count update steps, i.e. training batches, not epochs. Steps are numbered from 1; the position after step t is retained when t = burn_in + k · thinning for some k ≥ 1. Until the first position is retained, has_posterior is False and posterior refuses, so a posterior stage skips validation for the epochs of a burn-in.

SG-MCMC declares no early-stopping support, so a PosteriorStage refuses an early stopper for it and writes only a finalized checkpoint, holding the stacked retained samples.

Attributes:

Name	Type	Description
step_size	float	The Langevin step size ε, positive and finite.
burn_in	int	The number of initial update steps whose positions are discarded. Defaults to 0.
thinning	int	Retain one position every thinning steps after burn-in. Defaults to 1, which retains every position.
decay	float	The RMSprop decay of the squared-gradient average, in [0, 1). Defaults to 0.99.
stability	float	Added to √V before inverting it, positive. Defaults to 1e-5.

supports_early_stopping property

supports_early_stopping: bool

False: an early stopper would truncate the chain.

has_posterior property

has_posterior: bool

Whether at least one sample has been retained.

init

init(problem: PosteriorProblem) -> None

Start a chain at the problem's warm start, discarding any samples.

Parameters:

Name	Type	Description	Default
problem	PosteriorProblem	The posterior problem.	required

Source code in src/stagereg/jax/psgld.py

def init(self, problem: PosteriorProblem) -> None:
    """Start a chain at the problem's warm start, discarding any samples.

    Args:
        problem: The posterior problem.
    """
    self._problem = problem
    self._step = _make_step(
        problem,
        step_size=self.step_size,
        decay=self.decay,
        stability=self.stability,
    )
    self._parameters = problem.initial_parameters
    self._second_moment = jax.tree.map(jnp.zeros_like, problem.initial_parameters)
    self._steps = 0
    self._samples = []

update

update(batch: Batch, key: Array) -> Mapping[str, Any]

Take one Langevin step, retaining the new position if it is due.

Parameters:

Name	Type	Description	Default
batch	Batch	The training batch.	required
key	Array	The PRNG key of this step's injected noise.	required

Returns:

Type	Description
Mapping[str, Any]	{"loss": U}, the negative log posterior estimate at the
Mapping[str, Any]	position before the step.

Raises:

Type	Description
ValueError	If the method was not initialized.

Source code in src/stagereg/jax/psgld.py

def update(self, batch: Batch, key: jax.Array) -> Mapping[str, Any]:
    """Take one Langevin step, retaining the new position if it is due.

    Args:
        batch: The training batch.
        key: The PRNG key of this step's injected noise.

    Returns:
        ``{"loss": U}``, the negative log posterior estimate at the
        position before the step.

    Raises:
        ValueError: If the method was not initialized.
    """
    if self._step is None:
        raise ValueError("call init before update.")
    self._parameters, self._second_moment, loss = self._step(
        self._parameters,
        self._second_moment,
        batch.inputs,
        batch.targets,
        batch.sample_weight,
        key,
    )
    self._steps += 1
    retained_after = self._steps - self.burn_in
    if retained_after > 0 and retained_after % self.thinning == 0:
        self._samples.append(self._parameters)
    return {"loss": loss}

posterior

posterior() -> RetainedSamplesPosterior

Return the posterior of the samples retained so far.

Returns:

Type	Description
RetainedSamplesPosterior	A finite posterior with one draw per retained sample.

Raises:

Type	Description
ValueError	If no sample has been retained yet, e.g. during burn-in.

Source code in src/stagereg/jax/psgld.py

def posterior(self) -> RetainedSamplesPosterior:
    """Return the posterior of the samples retained so far.

    Returns:
        A finite posterior with one draw per retained sample.

    Raises:
        ValueError: If no sample has been retained yet, e.g. during
            burn-in.
    """
    return self._posterior_of(self.posterior_state())

state

state() -> PyTree

Return the chain's full state, from which it can resume.

Returns:

Type	Description
PyTree	The position, the preconditioner's squared-gradient average, the
PyTree	step count and the retained samples (stacked, possibly none).

Source code in src/stagereg/jax/psgld.py

def state(self) -> PyTree:
    """Return the chain's full state, from which it can resume.

    Returns:
        The position, the preconditioner's squared-gradient average, the
        step count and the retained samples (stacked, possibly none).
    """
    return {
        "parameters": self._parameters,
        "second_moment": self._second_moment,
        "steps": jnp.asarray(self._steps),
        "samples": _stack(self._samples, like=self._parameters),
    }

load_state

load_state(state: PyTree) -> None

Resume the chain from a state returned by state().

Parameters:

Name	Type	Description	Default
state	PyTree	A state previously returned by state.	required

Source code in src/stagereg/jax/psgld.py

def load_state(self, state: PyTree) -> None:
    """Resume the chain from a state returned by ``state()``.

    Args:
        state: A state previously returned by
            [`state`][stagereg.jax.PreconditionedSGLD.state].
    """
    self._parameters = state["parameters"]
    self._second_moment = state["second_moment"]
    self._steps = int(state["steps"])
    self._samples = _unstack(state["samples"])

posterior_state

posterior_state() -> PyTree

Return the retained samples, stacked along a leading axis.

Returns:

Type	Description
PyTree	The parameter tree with a leading sample axis.

Raises:

Type	Description
ValueError	If no sample has been retained yet.

Source code in src/stagereg/jax/psgld.py

def posterior_state(self) -> PyTree:
    """Return the retained samples, stacked along a leading axis.

    Returns:
        The parameter tree with a leading sample axis.

    Raises:
        ValueError: If no sample has been retained yet.
    """
    if not self._samples:
        raise ValueError(
            "no sample has been retained yet; train beyond burn_in "
            f"({self.burn_in} steps) plus thinning ({self.thinning} steps)."
        )
    return _stack(self._samples, like=self._parameters)

load_posterior

load_posterior(state: PyTree) -> None

Rebuild the posterior of the given retained samples after init.

Parameters:

Name	Type	Description	Default
state	PyTree	Stacked samples previously returned by posterior_state.	required

Source code in src/stagereg/jax/psgld.py

def load_posterior(self, state: PyTree) -> None:
    """Rebuild the posterior of the given retained samples after ``init``.

    Args:
        state: Stacked samples previously returned by
            [`posterior_state`][stagereg.jax.PreconditionedSGLD.posterior_state].
    """
    self._samples = _unstack(state)

RetainedSamplesPosterior dataclass

RetainedSamplesPosterior(
    mean_function: Callable[[PyTree, PyTree], Array],
    samples: PyTree,
)

A finite posterior whose draws are retained parameter samples.

Draw s is the posterior network with the s-th retained parameters, so it is the same mean function at every input, and every call uses all the draws.

Attributes:

Name	Type	Description
mean_function	Callable[[PyTree, PyTree], Array]	Maps (parameters, inputs) to the posterior network's mean predictions at inputs.
samples	PyTree	The retained parameter trees, stacked along a leading axis of length S.

num_draws property

num_draws: int

The number of retained samples S.

sample_means

sample_means(
    inputs: PyTree,
    key: Array,
    num_samples: int | None = None,
) -> jax.Array

Evaluate every retained sample's mean function at inputs.

Parameters:

Name	Type	Description	Default
inputs	PyTree	Model inputs with a leading batch dimension.	required
key	Array	Ignored: the draws are fixed.	required
num_samples	int | None	Must be None; all retained samples are used.	None

Returns:

Type	Description
Array	The mean predictions, shaped [S, *batch].

Raises:

Type	Description
ValueError	If num_samples is given.

Source code in src/stagereg/jax/psgld.py

def sample_means(
    self,
    inputs: PyTree,
    key: jax.Array,
    num_samples: int | None = None,
) -> jax.Array:
    """Evaluate every retained sample's mean function at ``inputs``.

    Args:
        inputs: Model inputs with a leading batch dimension.
        key: Ignored: the draws are fixed.
        num_samples: Must be ``None``; all retained samples are used.

    Returns:
        The mean predictions, shaped ``[S, *batch]``.

    Raises:
        ValueError: If ``num_samples`` is given.
    """
    del key
    if num_samples is not None:
        raise ValueError(
            "a finite posterior uses all its retained samples; "
            "num_samples must be None."
        )
    return jax.vmap(self.mean_function, in_axes=(0, None))(self.samples, inputs)

stagereg.jax.rng

split_key

split_key(key: Array) -> tuple[jax.Array, jax.Array]

Return the next persistent key and one key for a single operation.

Parameters:

Name	Type	Description	Default
key	Array	The current JAX PRNG key.	required

Returns:

Type	Description
Array	A tuple (next_key, operation_key) where next_key should
Array	replace key for subsequent splits and operation_key is
tuple[Array, Array]	consumed by a single random operation.

Source code in src/stagereg/jax/rng.py

def split_key(key: jax.Array) -> tuple[jax.Array, jax.Array]:
    """Return the next persistent key and one key for a single operation.

    Args:
        key: The current JAX PRNG key.

    Returns:
        A tuple ``(next_key, operation_key)`` where ``next_key`` should
        replace ``key`` for subsequent splits and ``operation_key`` is
        consumed by a single random operation.
    """
    next_key, operation_key = jax.random.split(key)
    return next_key, operation_key

stagereg.jax.state

NnxSnapshot dataclass

NnxSnapshot(graphdef: GraphDef, state: State)

An opaque, complete snapshot of an NNX graph and its state.

Attributes:

Name	Type	Description
graphdef	GraphDef	The static NNX graph definition.
state	State	The NNX variable state captured at snapshot time.

create_optimizer

create_optimizer(
    model: Module, transformation: GradientTransformation
) -> nnx.Optimizer

Create an NNX optimizer that updates all model parameters.

Parameters:

Name	Type	Description	Default
model	Module	The NNX module whose parameters will be optimized.	required
transformation	GradientTransformation	The Optax gradient transformation defining the optimization algorithm.	required

Returns:

Type	Description
Optimizer	An nnx.Optimizer bound to model's parameters.

Source code in src/stagereg/jax/state.py

def create_optimizer(
    model: nnx.Module, transformation: optax.GradientTransformation
) -> nnx.Optimizer:
    """Create an NNX optimizer that updates all model parameters.

    Args:
        model: The NNX module whose parameters will be optimized.
        transformation: The Optax gradient transformation defining the
            optimization algorithm.

    Returns:
        An ``nnx.Optimizer`` bound to ``model``'s parameters.
    """
    return nnx.Optimizer(model, transformation, wrt=nnx.Param)

initialize_training_state

initialize_training_state(
    model: Module,
    optimizer: Optimizer,
    *,
    rng_key: Array,
    model_name: str = "model",
    optimizer_name: str = "optimizer",
) -> TrainingState

Create generic training state backed by an NNX model and optimizer.

Parameters:

Name	Type	Description	Default
model	Module	The NNX module to register as a model component.	required
optimizer	Optimizer	The NNX optimizer to register as an optimizer state.	required
rng_key	Array	The initial JAX PRNG key for the training state.	required
model_name	str	Key under which model is stored in TrainingState.model_components. Defaults to "model".	'model'
optimizer_name	str	Key under which optimizer is stored in TrainingState.optimizer_states. Defaults to "optimizer".	'optimizer'

Returns:

Type	Description
TrainingState	A TrainingState with model and
TrainingState	optimizer registered and rng_state initialized to
TrainingState	rng_key.

Source code in src/stagereg/jax/state.py

def initialize_training_state(
    model: nnx.Module,
    optimizer: nnx.Optimizer,
    *,
    rng_key: jax.Array,
    model_name: str = "model",
    optimizer_name: str = "optimizer",
) -> TrainingState:
    """Create generic training state backed by an NNX model and optimizer.

    Args:
        model: The NNX module to register as a model component.
        optimizer: The NNX optimizer to register as an optimizer state.
        rng_key: The initial JAX PRNG key for the training state.
        model_name: Key under which ``model`` is stored in
            ``TrainingState.model_components``. Defaults to ``"model"``.
        optimizer_name: Key under which ``optimizer`` is stored in
            ``TrainingState.optimizer_states``. Defaults to
            ``"optimizer"``.

    Returns:
        A [`TrainingState`][stagereg.core.TrainingState] with ``model`` and
        ``optimizer`` registered and ``rng_state`` initialized to
        ``rng_key``.
    """
    return TrainingState(
        model_components={model_name: model},
        optimizer_states={optimizer_name: optimizer},
        rng_state=rng_key,
    )

freeze_training_state

freeze_training_state(
    state: TrainingState,
) -> TrainingState

Create an independent, checkpoint-safe snapshot of a training state.

TrainingState holds several mutable containers (metric_history, checkpoint_registry, parameter_roles) that training runners keep appending to and mutating in place across epochs. A plain dataclasses.replace only shallow-copies the state, so a checkpoint built that way would keep sharing those containers with live training and silently drift after being "saved". This function deep-copies every mutable container field instead.

The live model_components and optimizer_states are dropped from the returned snapshot rather than copied, since they hold mutable NNX modules/optimizers that keep changing across epochs; callers that need an immutable capture of model or optimizer state should use snapshot instead and store it alongside the frozen training state (see Checkpoint).

Parameters:

Name	Type	Description	Default
state	TrainingState	The live training state to snapshot.	required

Returns:

Type	Description
TrainingState	A new TrainingState whose mutable
TrainingState	container fields are independent copies, unaffected by later
TrainingState	in-place mutation of state, and whose
TrainingState	model_components/optimizer_states are empty.

Source code in src/stagereg/jax/state.py

def freeze_training_state(state: TrainingState) -> TrainingState:
    """Create an independent, checkpoint-safe snapshot of a training state.

    ``TrainingState`` holds several mutable containers (``metric_history``,
    ``checkpoint_registry``, ``parameter_roles``) that training runners
    keep appending to and mutating in place across epochs. A plain
    ``dataclasses.replace`` only shallow-copies the state, so a checkpoint
    built that way would keep sharing those containers with live
    training and silently drift after being "saved". This function
    deep-copies every mutable container field instead.

    The live ``model_components`` and ``optimizer_states`` are dropped
    from the returned snapshot rather than copied, since they hold
    mutable NNX modules/optimizers that keep changing across epochs;
    callers that need an immutable capture of model or optimizer state
    should use [`snapshot`][stagereg.jax.snapshot] instead and store it
    alongside the frozen training state (see
    [`Checkpoint`][stagereg.core.Checkpoint]).

    Args:
        state: The live training state to snapshot.

    Returns:
        A new [`TrainingState`][stagereg.core.TrainingState] whose mutable
        container fields are independent copies, unaffected by later
        in-place mutation of ``state``, and whose
        ``model_components``/``optimizer_states`` are empty.
    """
    return replace(
        state,
        model_components={},
        optimizer_states={},
        parameter_roles=dict(state.parameter_roles),
        checkpoint_registry=dict(state.checkpoint_registry),
        metric_history={
            name: list(values) for name, values in state.metric_history.items()
        },
    )

snapshot

snapshot(module: Module) -> NnxSnapshot

Capture an NNX module's graph definition and complete variable state.

Parameters:

Name	Type	Description	Default
module	Module	The NNX module to snapshot.	required

Returns:

Type	Description
NnxSnapshot	An NnxSnapshot containing module's
NnxSnapshot	graph definition and variable state at the time of the call.

Source code in src/stagereg/jax/state.py

def snapshot(module: nnx.Module) -> NnxSnapshot:
    """Capture an NNX module's graph definition and complete variable state.

    Args:
        module: The NNX module to snapshot.

    Returns:
        An [`NnxSnapshot`][stagereg.jax.NnxSnapshot] containing ``module``'s
        graph definition and variable state at the time of the call.
    """
    graphdef, state = nnx.split(module)
    copied_state = jax.tree.map(_copy_snapshot_leaf, state)
    return NnxSnapshot(graphdef=graphdef, state=copied_state)

restore_checkpoint

restore_checkpoint(
    checkpoint: Checkpoint,
    *,
    state: TrainingState,
    model: Module,
    optimizer: Optimizer,
    model_name: str = "model",
    optimizer_name: str = "optimizer",
) -> None

Restore a JAX checkpoint into existing live training objects.

The model and optimizer are updated in place so callers retaining their identities, including stage registries, continue to reference the restored objects. Checkpoint-safe training states omit mutable model and optimizer objects, so this helper re-registers the supplied live objects after restoring the backend-neutral state fields.

Parameters:

Name	Type	Description	Default
checkpoint	Checkpoint	Checkpoint containing NNX model and optimizer snapshots.	required
state	TrainingState	Existing shared training state to restore in place.	required
model	Module	Live model object updated from checkpoint.parameters.	required
optimizer	Optimizer	Live optimizer updated from checkpoint.optimizer_state.	required
model_name	str	Registry name for model. Defaults to "model".	'model'
optimizer_name	str	Registry name for optimizer. Defaults to "optimizer".	'optimizer'

Raises:

Type	Description
TypeError	If the checkpoint does not contain NNX snapshots or a JAX random key.
ValueError	If the live model or optimizer is incompatible with the checkpoint snapshot.

Source code in src/stagereg/jax/state.py

def restore_checkpoint(
    checkpoint: Checkpoint,
    *,
    state: TrainingState,
    model: nnx.Module,
    optimizer: nnx.Optimizer,
    model_name: str = "model",
    optimizer_name: str = "optimizer",
) -> None:
    """Restore a JAX checkpoint into existing live training objects.

    The model and optimizer are updated in place so callers retaining their
    identities, including stage registries, continue to reference the restored
    objects. Checkpoint-safe training states omit mutable model and optimizer
    objects, so this helper re-registers the supplied live objects after
    restoring the backend-neutral state fields.

    Args:
        checkpoint: Checkpoint containing NNX model and optimizer snapshots.
        state: Existing shared training state to restore in place.
        model: Live model object updated from ``checkpoint.parameters``.
        optimizer: Live optimizer updated from ``checkpoint.optimizer_state``.
        model_name: Registry name for ``model``. Defaults to ``"model"``.
        optimizer_name: Registry name for ``optimizer``. Defaults to
            ``"optimizer"``.

    Raises:
        TypeError: If the checkpoint does not contain NNX snapshots or a JAX
            random key.
        ValueError: If the live model or optimizer is incompatible with the
            checkpoint snapshot.
    """
    if not isinstance(checkpoint.parameters, NnxSnapshot):
        raise TypeError("checkpoint.parameters must be an NnxSnapshot.")
    if not isinstance(checkpoint.optimizer_state, NnxSnapshot):
        raise TypeError("checkpoint.optimizer_state must be an NnxSnapshot.")
    if not isinstance(checkpoint.rng_state, jax.Array):
        raise TypeError("checkpoint.rng_state must be a JAX random key.")

    _validate_graph(model, checkpoint.parameters, kind="model")
    _validate_state_structure(
        optimizer,
        checkpoint.optimizer_state,
        kind="optimizer",
    )
    restored_optimizer = nnx.merge(
        checkpoint.optimizer_state.graphdef,
        checkpoint.optimizer_state.state,
    )
    nnx.update(model, checkpoint.parameters.state)
    _restore_optimizer_static_fields(optimizer, restored_optimizer)
    nnx.update(optimizer, checkpoint.optimizer_state.state)
    _restore_training_state(state, checkpoint.state)
    state.rng_state = checkpoint.rng_state
    state.register_component(model_name, model)
    state.register_optimizer(optimizer_name, optimizer)

stagereg.jax.supervised

make_train_step

make_train_step(
    loss: SupervisedLoss,
    *,
    metrics: Sequence[BatchMetricSpec] = (),
) -> Callable[..., Mapping[str, jax.Array]]

Create a JIT-compiled NNX/Optax supervised training step.

Parameters:

Name	Type	Description	Default
loss	SupervisedLoss	A callable computing the supervised loss given a model, inputs, targets, sample weights, a PRNG key, and a training flag.	required

Returns:

Type	Description
Callable[..., Mapping[str, Array]]	A JIT-compiled function ``train_step(model, optimizer, inputs,
Callable[..., Mapping[str, Array]]	targets, sample_weight, key)`` that performs one gradient update
Callable[..., Mapping[str, Array]]	in place and returns a mapping containing "loss" plus
Callable[..., Mapping[str, Array]]	registered batch metric values. Loss is evaluated on the
Callable[..., Mapping[str, Array]]	pre-update training-state model, while batch metrics are computed
Callable[..., Mapping[str, Array]]	on the same pre-update parameters with training=False to
Callable[..., Mapping[str, Array]]	avoid a second state-mutating training-mode forward pass.

Source code in src/stagereg/jax/supervised.py

def make_train_step(
    loss: SupervisedLoss,
    *,
    metrics: Sequence[BatchMetricSpec] = (),
) -> Callable[..., Mapping[str, jax.Array]]:
    """Create a JIT-compiled NNX/Optax supervised training step.

    Args:
        loss: A callable computing the supervised loss given a model,
            inputs, targets, sample weights, a PRNG key, and a
            ``training`` flag.

    Returns:
        A JIT-compiled function ``train_step(model, optimizer, inputs,
        targets, sample_weight, key)`` that performs one gradient update
        in place and returns a mapping containing ``"loss"`` plus
        registered batch metric values. Loss is evaluated on the
        pre-update training-state model, while batch metrics are computed
        on the same pre-update parameters with ``training=False`` to
        avoid a second state-mutating training-mode forward pass.
    """

    @nnx.jit
    def train_step(
        model: nnx.Module,
        optimizer: nnx.Optimizer,
        inputs: PyTree,
        targets: jax.Array,
        sample_weight: jax.Array | None,
        key: jax.Array,
    ) -> Mapping[str, jax.Array]:
        def loss_fn(current_model: nnx.Module) -> jax.Array:
            return loss(
                current_model,
                inputs,
                targets,
                sample_weight,
                key,
                True,
            )

        loss_value, gradients = nnx.value_and_grad(loss_fn)(model)
        values: dict[str, jax.Array] = {"loss": loss_value}
        for spec in metrics:
            values[spec.name] = spec.metric(
                model,
                inputs,
                targets,
                sample_weight,
                key,
                False,
            )
        optimizer.update(model, gradients)
        return values

    return train_step

run_supervised

run_supervised(
    *,
    model: Module,
    optimizer: Optimizer,
    train_loader: LoaderFactory,
    loss: SupervisedLoss,
    state: TrainingState,
    epochs: int,
    validation: ValidationStrategy | None = None,
    early_stopper: EarlyStopper | None = None,
    event_sinks: Sequence[EventSink] = (),
    checkpoint_store: CheckpointStore | None = None,
    checkpoint_key: str | None = None,
    stage: str = "supervised",
    model_name: str = "model",
    optimizer_name: str = "optimizer",
    metrics: MetricSuite | None = None,
) -> StageResult

Train a single NNX model for a fixed or early-stopped number of epochs.

The runner does not restore the best checkpoint: when training ends, model holds the parameters of the last epoch run. To continue from the best model, load the checkpoint from checkpoint_store and pass it to restore_checkpoint.

Parameters:

Name	Type	Description	Default
model	Module	The NNX module to train, mutated in place.	required
optimizer	Optimizer	The NNX optimizer used to update model, mutated in place.	required
train_loader	LoaderFactory	Factory producing the training batch loader for a given split and epoch.	required
loss	SupervisedLoss	A callable computing the supervised loss given a model, inputs, targets, sample weights, a PRNG key, and a training flag.	required
state	TrainingState	The training state to update in place across epochs.	required
epochs	int	The maximum number of epochs to run.	required
validation	ValidationStrategy | None	An optional strategy used to evaluate state after each epoch. Required if early_stopper monitors a validation metric.	None
early_stopper	EarlyStopper | None	An optional policy that stops training early based on a monitored training or validation metric.	None
event_sinks	Sequence[EventSink]	Sinks notified of epoch, validation, best-model, and early-stop events.	()
checkpoint_store	CheckpointStore | None	An optional store used to persist the best checkpoint when early_stopper reports an improvement.	None
checkpoint_key	str | None	The checkpoint key under which the best checkpoint is saved. Defaults to None, which resolves to the stage-scoped key f"{stage}/best" so that stages sharing one checkpoint_store never overwrite each other; the MeanStage and GammaVarianceStage defaults resolve the same way.	None
stage	str	The stage name recorded on state and emitted events, and the stage segment of every metric tag recorded in state.metric_history. Defaults to "supervised".	'supervised'
model_name	str	Name under which model is registered in state. Defaults to "model".	'model'
optimizer_name	str	Name under which optimizer is registered in state. Defaults to "optimizer".	'optimizer'
metrics	MetricSuite | None	Optional registered batch/epoch metrics for training. When omitted, only loss is collected.	None

Returns:

Type	Description
StageResult	A StageResult with the final
StageResult	state, the last recorded training metrics under their bare
StageResult	names, and the final training loss. state.metric_history
StageResult	records every metric under its metric tag stage/split/metric.

Raises:

Type	Description
ValueError	If epochs is not positive, if stage is empty or contains /, if early_stopper monitors a validation metric without a validation strategy, or if the monitored metric is not produced by training or validation.
TypeError	If state.rng_state is not a JAX random key.

Source code in src/stagereg/jax/supervised.py

def run_supervised(
    *,
    model: nnx.Module,
    optimizer: nnx.Optimizer,
    train_loader: LoaderFactory,
    loss: SupervisedLoss,
    state: TrainingState,
    epochs: int,
    validation: ValidationStrategy | None = None,
    early_stopper: EarlyStopper | None = None,
    event_sinks: Sequence[EventSink] = (),
    checkpoint_store: CheckpointStore | None = None,
    checkpoint_key: str | None = None,
    stage: str = "supervised",
    model_name: str = "model",
    optimizer_name: str = "optimizer",
    metrics: MetricSuite | None = None,
) -> StageResult:
    """Train a single NNX model for a fixed or early-stopped number of epochs.

    The runner does not restore the best checkpoint: when training ends,
    ``model`` holds the parameters of the last epoch run. To continue from the
    best model, load the checkpoint from ``checkpoint_store`` and pass it to
    [`restore_checkpoint`][stagereg.jax.restore_checkpoint].

    Args:
        model: The NNX module to train, mutated in place.
        optimizer: The NNX optimizer used to update ``model``, mutated
            in place.
        train_loader: Factory producing the training batch loader for a
            given split and epoch.
        loss: A callable computing the supervised loss given a model,
            inputs, targets, sample weights, a PRNG key, and a
            ``training`` flag.
        state: The training state to update in place across epochs.
        epochs: The maximum number of epochs to run.
        validation: An optional strategy used to evaluate ``state``
            after each epoch. Required if ``early_stopper`` monitors a
            validation metric.
        early_stopper: An optional policy that stops training early
            based on a monitored training or validation metric.
        event_sinks: Sinks notified of epoch, validation, best-model,
            and early-stop events.
        checkpoint_store: An optional store used to persist the best
            checkpoint when ``early_stopper`` reports an improvement.
        checkpoint_key: The checkpoint key under which the best checkpoint
            is saved. Defaults to ``None``, which resolves to the stage-scoped
            key ``f"{stage}/best"`` so that stages sharing one
            ``checkpoint_store`` never overwrite each other; the
            [`MeanStage`][stagereg.jax.MeanStage] and
            [`GammaVarianceStage`][stagereg.jax.GammaVarianceStage] defaults
            resolve the same way.
        stage: The stage name recorded on ``state`` and emitted events,
            and the stage segment of every metric tag recorded in
            ``state.metric_history``. Defaults to ``"supervised"``.
        model_name: Name under which ``model`` is registered in ``state``.
            Defaults to ``"model"``.
        optimizer_name: Name under which ``optimizer`` is registered in
            ``state``. Defaults to ``"optimizer"``.
        metrics: Optional registered batch/epoch metrics for training. When
            omitted, only loss is collected.

    Returns:
        A [`StageResult`][stagereg.core.StageResult] with the final
        ``state``, the last recorded training metrics under their bare
        names, and the final training loss. ``state.metric_history``
        records every metric under its metric tag ``stage/split/metric``.

    Raises:
        ValueError: If ``epochs`` is not positive, if ``stage`` is empty
            or contains ``/``, if ``early_stopper`` monitors a validation
            metric without a ``validation`` strategy, or if the monitored
            metric is not produced by training or validation.
        TypeError: If ``state.rng_state`` is not a JAX random key.
    """
    check_epoch_loop_arguments(
        state=state,
        epochs=epochs,
        stage=stage,
        validation=validation,
        early_stopper=early_stopper,
    )
    state.register_component(model_name, model)
    state.register_optimizer(optimizer_name, optimizer)
    metric_suite = metrics if metrics is not None else MetricSuite()
    train_step = make_train_step(loss, metrics=metric_suite.batch)

    def step(batch: Batch, key: jax.Array, /) -> Mapping[str, jax.Array]:
        return train_step(
            model,
            optimizer,
            batch.inputs,
            batch.targets,
            batch.sample_weight,
            key,
        )

    def snapshot_state() -> StateSnapshot:
        return StateSnapshot(
            parameters=snapshot(model), optimizer_state=snapshot(optimizer)
        )

    def epoch_predictions(batch: Batch, key: jax.Array, /) -> EpochPredictionData:
        return resolve_epoch_prediction_data(
            suite=metric_suite, model=model, batch=batch, key=key
        )

    return run_epoch_loop(
        step=step,
        snapshot_state=snapshot_state,
        train_loader=train_loader,
        state=state,
        epochs=epochs,
        stage=stage,
        checkpoint_key=checkpoint_key,
        validation=validation,
        early_stopper=early_stopper,
        event_sinks=event_sinks,
        checkpoint_store=checkpoint_store,
        metrics=metric_suite,
        epoch_predictions=epoch_predictions,
    )

stagereg.jax.supervised_staged

materialize_residual_loader

materialize_residual_loader(
    mean_model: Module,
    source_loader: LoaderFactory,
    *,
    splits: Sequence[str] = ("train", "validation"),
    source_epoch: int = 0,
) -> LoaderFactory

Materialize detached squared residual targets for fixed data splits.

The source loader is consumed exactly once per configured split. The returned loader replays the cached batches for every later epoch, making the Step 1 mean predictions and source data a fixed Step 2 training snapshot.

Parameters:

Name	Type	Description	Default
mean_model	Module	Trained deterministic mean model.	required
source_loader	LoaderFactory	Loader providing inputs and original regression targets.	required
splits	Sequence[str]	Split names to materialize. Defaults to training and validation.	('train', 'validation')
source_epoch	int	Source-loader epoch used for the one-time snapshot.	0

Returns:

Type	Description
LoaderFactory	A loader factory serving cached batches whose targets are detached
LoaderFactory	squared residuals.

Raises:

Type	Description
ValueError	If splits are empty, duplicated, unknown at replay time, or produce no batches; if a source batch lacks targets; or if mean predictions and targets have different shapes.

Source code in src/stagereg/jax/supervised_staged.py

def materialize_residual_loader(
    mean_model: nnx.Module,
    source_loader: LoaderFactory,
    *,
    splits: Sequence[str] = ("train", "validation"),
    source_epoch: int = 0,
) -> LoaderFactory:
    """Materialize detached squared residual targets for fixed data splits.

    The source loader is consumed exactly once per configured split. The
    returned loader replays the cached batches for every later epoch, making
    the Step 1 mean predictions and source data a fixed Step 2 training
    snapshot.

    Args:
        mean_model: Trained deterministic mean model.
        source_loader: Loader providing inputs and original regression targets.
        splits: Split names to materialize. Defaults to training and validation.
        source_epoch: Source-loader epoch used for the one-time snapshot.

    Returns:
        A loader factory serving cached batches whose targets are detached
        squared residuals.

    Raises:
        ValueError: If splits are empty, duplicated, unknown at replay time, or
            produce no batches; if a source batch lacks targets; or if mean
            predictions and targets have different shapes.
    """
    split_names = tuple(splits)
    if not split_names:
        raise ValueError("splits must contain at least one split name.")
    if any(not name for name in split_names):
        raise ValueError("split names must not be empty.")
    if len(set(split_names)) != len(split_names):
        raise ValueError("split names must be unique.")

    prediction_model = nnx.clone(mean_model)
    prediction_model.eval()
    cached: dict[str, tuple[Batch, ...]] = {}

    for split in split_names:
        residual_batches = tuple(
            _materialize_residual_batch(prediction_model, batch)
            for batch in source_loader(split=split, epoch=source_epoch)
        )
        if not residual_batches:
            raise ValueError(f"source split {split!r} must provide at least one batch.")
        cached[split] = residual_batches

    def residual_loader(*, split: str, epoch: int) -> tuple[Batch, ...]:
        del epoch
        if split not in cached:
            raise ValueError(f"residual split {split!r} was not materialized.")
        return cached[split]

    return residual_loader

SupervisedStageOptions dataclass

SupervisedStageOptions(
    epochs: int,
    validation: ValidationStrategy | None = None,
    early_stopper: EarlyStopper | None = None,
    event_sinks: Sequence[EventSink] = (),
    checkpoint_store: CheckpointStore | None = None,
    checkpoint_key: str | None = None,
    metrics: MetricSuite = MetricSuite(),
)

Shared epoch-runner options for concrete supervised stages.

Attributes:

Name	Type	Description
epochs	int	Maximum number of training epochs.
validation	ValidationStrategy | None	Optional validation strategy.
early_stopper	EarlyStopper | None	Optional early-stopping policy.
event_sinks	Sequence[EventSink]	Event consumers notified by the runner.
checkpoint_store	CheckpointStore | None	Optional best-checkpoint store.
checkpoint_key	str | None	The checkpoint key the stage saves its best checkpoint under. Defaults to None, which resolves to the stage-scoped key f"{stage}/best" (mean/best, variance/best), the same default as run_supervised, so stages sharing one checkpoint_store never overwrite each other's checkpoints.
metrics	MetricSuite	Batch and epoch metric registrations.

MeanStage dataclass

MeanStage(
    model: Module,
    optimizer: Optimizer,
    train_loader: LoaderFactory,
    options: SupervisedStageOptions,
    model_name: str = "mean_model",
    optimizer_name: str = "mean_optimizer",
    loss: SupervisedLoss = (
        lambda: make_supervised_loss(SquaredErrorLoss())
    )(),
)

Concrete Step 1 stage training a deterministic mean model with MSE.

Attributes:

Name	Type	Description
model	Module	NNX model producing deterministic mean predictions.
optimizer	Optimizer	NNX optimizer bound to model.
train_loader	LoaderFactory	Factory producing original regression batches.
options	SupervisedStageOptions	Shared supervised-runner options.
model_name	str	State registry name for the mean model.
optimizer_name	str	State registry name for the mean optimizer.
loss	SupervisedLoss	Scalar supervised loss used to train the mean model.

prepare

prepare(state: TrainingState) -> None

Register Step 1 ownership and initialize the staged lifecycle.

Parameters:

Name	Type	Description	Default
state	TrainingState	Shared staged training state.	required

Raises:

Type	Description
ValueError	If the lifecycle or component ownership is invalid.

Source code in src/stagereg/jax/supervised_staged.py

def prepare(self, state: TrainingState) -> None:
    """Register Step 1 ownership and initialize the staged lifecycle.

    Args:
        state: Shared staged training state.

    Raises:
        ValueError: If the lifecycle or component ownership is invalid.
    """
    if state.lifecycle_state not in (StageState.NEW, StageState.INITIALIZED):
        raise ValueError("mean stage requires NEW or INITIALIZED lifecycle state.")
    _validate_named_registration(
        state.model_components,
        self.model_name,
        self.model,
        kind="model component",
    )
    _validate_named_registration(
        state.optimizer_states,
        self.optimizer_name,
        self.optimizer,
        kind="optimizer state",
    )
    _validate_parameter_role(state, self.model_name, ParameterRole.MEAN)
    state.register_component(self.model_name, self.model)
    state.register_optimizer(self.optimizer_name, self.optimizer)
    state.parameter_roles[self.model_name] = ParameterRole.MEAN
    state.lifecycle_state = StageState.INITIALIZED
    state.active_stage = self.name

train

train(state: TrainingState) -> StageResult

Train the mean model and transition the workflow to MEAN_READY.

With an early stopper and a checkpoint store configured, the stage then restores its best checkpoint, so the live mean model holds the best epoch's weights, and saves it again under the same key as a finalized MEAN_READY checkpoint with metadata {"stage": "mean", "stage_complete": True}.

Parameters:

Name	Type	Description	Default
state	TrainingState	Prepared staged training state.	required

Returns:

Type	Description
StageResult	The supervised runner result, or the restored epoch's training
StageResult	metrics when the best checkpoint was restored.

Raises:

Type	Description
ValueError	If the stage was not prepared or training is non-finite.

Source code in src/stagereg/jax/supervised_staged.py

def train(self, state: TrainingState) -> StageResult:
    """Train the mean model and transition the workflow to ``MEAN_READY``.

    With an early stopper and a checkpoint store configured, the stage then
    restores its best checkpoint, so the live mean model holds the best
    epoch's weights, and saves it again under the same key as a finalized
    ``MEAN_READY`` checkpoint with metadata
    ``{"stage": "mean", "stage_complete": True}``.

    Args:
        state: Prepared staged training state.

    Returns:
        The supervised runner result, or the restored epoch's training
        metrics when the best checkpoint was restored.

    Raises:
        ValueError: If the stage was not prepared or training is non-finite.
    """
    if state.lifecycle_state is not StageState.INITIALIZED:
        raise ValueError("mean stage must be prepared before training.")
    return _train_stage(
        self,
        state,
        train_loader=self.train_loader,
        validation=self.options.validation,
        ready=StageState.MEAN_READY,
    )

restore

restore(
    state: TrainingState, checkpoint: Checkpoint
) -> None

Restore the mean stage's finalized checkpoint into state.

The live model and optimizer take the checkpoint's weights in place and are registered under model_name and optimizer_name. The restore is clean-slate, like restore_checkpoint: other registered components and optimizers are dropped. Afterwards the state passes validate, so the variance stage can be prepared or restored next.

Parameters:

Name	Type	Description	Default
state	TrainingState	Training state to restore into, typically a fresh one.	required
checkpoint	Checkpoint	The mean stage's finalized checkpoint.	required

Raises:

Type	Description
ValueError	If checkpoint is not finalized by this stage, that is its lifecycle state is not MEAN_READY or its metadata lacks "stage": "mean" and "stage_complete": True, or if it was saved with another model_name; this is checked before state or the live objects change. Also raised, as by restore_checkpoint, if the live model or optimizer is incompatible with the checkpoint, which may leave them and state partially restored.
TypeError	As by restore_checkpoint, if the checkpoint does not hold NNX snapshots or a JAX random key.

Source code in src/stagereg/jax/supervised_staged.py

def restore(self, state: TrainingState, checkpoint: Checkpoint) -> None:
    """Restore the mean stage's finalized checkpoint into ``state``.

    The live model and optimizer take the checkpoint's weights in place and
    are registered under ``model_name`` and ``optimizer_name``. The restore
    is clean-slate, like
    [`restore_checkpoint`][stagereg.jax.restore_checkpoint]: other registered
    components and optimizers are dropped. Afterwards the state passes
    [`validate`][stagereg.jax.MeanStage.validate], so the variance stage can
    be prepared or restored next.

    Args:
        state: Training state to restore into, typically a fresh one.
        checkpoint: The mean stage's finalized checkpoint.

    Raises:
        ValueError: If ``checkpoint`` is not finalized by this stage, that
            is its lifecycle state is not ``MEAN_READY`` or its metadata
            lacks ``"stage": "mean"`` and ``"stage_complete": True``, or
            if it was saved with another ``model_name``; this is checked
            before ``state`` or the live objects change. Also
            raised, as by
            [`restore_checkpoint`][stagereg.jax.restore_checkpoint], if the
            live model or optimizer is incompatible with the checkpoint,
            which may leave them and ``state`` partially restored.
        TypeError: As by
            [`restore_checkpoint`][stagereg.jax.restore_checkpoint], if the
            checkpoint does not hold NNX snapshots or a JAX random key.
    """
    require_finalized(checkpoint, stage=self.name, ready=StageState.MEAN_READY)
    require_component_names(
        checkpoint, model_name=self.model_name, role=ParameterRole.MEAN
    )
    self._restore_live(state, checkpoint)

validate

validate(state: TrainingState) -> ValidationResult

Validate mean-stage lifecycle and ownership invariants.

Parameters:

Name	Type	Description	Default
state	TrainingState	Shared staged training state.	required

Returns:

Type	Description
ValidationResult	A validation result describing whether Step 1 is ready.

Source code in src/stagereg/jax/supervised_staged.py

def validate(self, state: TrainingState) -> ValidationResult:
    """Validate mean-stage lifecycle and ownership invariants.

    Args:
        state: Shared staged training state.

    Returns:
        A validation result describing whether Step 1 is ready.
    """
    passed = (
        state.lifecycle_state is StageState.MEAN_READY
        and state.model_components.get(self.model_name) is self.model
        and state.optimizer_states.get(self.optimizer_name) is self.optimizer
        and state.parameter_roles.get(self.model_name) is ParameterRole.MEAN
    )
    return ValidationResult(
        passed=passed,
        message=None if passed else "mean stage invariants are not satisfied.",
    )

select_checkpoint

select_checkpoint(state: TrainingState) -> CheckpointRef

Return a reference to the mean stage's best checkpoint.

Parameters:

Name	Type	Description	Default
state	TrainingState	Shared staged training state.	required

Returns:

Type	Description
CheckpointRef	Reference to the checkpoint key the stage saves under: the
CheckpointRef	configured checkpoint_key, or f"{stage}/best" when none
CheckpointRef	was configured.

Raises:

Type	Description
ValueError	If no checkpoint exists under that key.

Source code in src/stagereg/jax/supervised_staged.py

def select_checkpoint(self, state: TrainingState) -> CheckpointRef:
    """Return a reference to the mean stage's best checkpoint.

    Args:
        state: Shared staged training state.

    Returns:
        Reference to the checkpoint key the stage saves under: the
        configured ``checkpoint_key``, or ``f"{stage}/best"`` when none
        was configured.

    Raises:
        ValueError: If no checkpoint exists under that key.
    """
    del state
    return _select_checkpoint(self)

GammaVarianceStage dataclass

GammaVarianceStage(
    model: Module,
    optimizer: Optimizer,
    source_loader: LoaderFactory,
    options: SupervisedStageOptions,
    mean_model_name: str = "mean_model",
    mean_optimizer_name: str = "mean_optimizer",
    model_name: str = "variance_model",
    optimizer_name: str = "variance_optimizer",
    splits: Sequence[str] = ("train", "validation"),
    source_epoch: int = 0,
    validation_factory: Callable[
        [LoaderFactory], ValidationStrategy
    ]
    | None = None,
    loss: SupervisedLoss = (
        lambda: make_supervised_loss(
            NegativeLogLikelihoodLoss(
                target_transform=add_epsilon()
            )
        )
    )(),
)

Concrete Step 2 stage fitting Gamma-distributed squared residuals.

Attributes:

Name	Type	Description
model	Module	NNX model producing Gamma residual predictions.
optimizer	Optimizer	NNX optimizer bound to model.
source_loader	LoaderFactory	Factory producing original regression batches.
options	SupervisedStageOptions	Shared supervised-runner options.
mean_model_name	str	Registry name of the prepared mean model.
mean_optimizer_name	str	Registry name of the mean optimizer, kept registered when the stage restores its best checkpoint.
model_name	str	State registry name for the variance model.
optimizer_name	str	State registry name for the variance optimizer.
splits	Sequence[str]	Source split names materialized as residual batches.
source_epoch	int	Source-loader epoch used for residual materialization.
validation_factory	Callable[[LoaderFactory], ValidationStrategy] | None	Optional factory adapting the residual loader into a validation strategy.
loss	SupervisedLoss	Scalar supervised loss used to train the variance model.

prepare

prepare(state: TrainingState) -> None

Validate Step 1 output and materialize fixed residual targets.

Parameters:

Name	Type	Description	Default
state	TrainingState	Shared state whose mean component is ready.	required

Raises:

Type	Description
TypeError	If the registered mean component is not an NNX module.
ValueError	If lifecycle, ownership, or registrations are invalid.

Source code in src/stagereg/jax/supervised_staged.py

def prepare(self, state: TrainingState) -> None:
    """Validate Step 1 output and materialize fixed residual targets.

    Args:
        state: Shared state whose mean component is ready.

    Raises:
        TypeError: If the registered mean component is not an NNX module.
        ValueError: If lifecycle, ownership, or registrations are invalid.
    """
    if state.lifecycle_state is not StageState.MEAN_READY:
        raise ValueError("variance stage requires MEAN_READY lifecycle state.")
    if self.mean_model_name not in state.model_components:
        raise ValueError(
            f"mean model component {self.mean_model_name!r} is not registered."
        )
    mean_model = state.model_components[self.mean_model_name]
    if not isinstance(mean_model, nnx.Module):
        raise TypeError("registered mean model must be an NNX module.")
    if state.parameter_roles.get(self.mean_model_name) is not ParameterRole.MEAN:
        raise ValueError("registered mean model must have the MEAN parameter role.")
    _validate_named_registration(
        state.model_components,
        self.model_name,
        self.model,
        kind="model component",
    )
    _validate_named_registration(
        state.optimizer_states,
        self.optimizer_name,
        self.optimizer,
        kind="optimizer state",
    )
    _validate_parameter_role(state, self.model_name, ParameterRole.VARIANCE)

    residual_loader = materialize_residual_loader(
        mean_model,
        self.source_loader,
        splits=self.splits,
        source_epoch=self.source_epoch,
    )
    validation = (
        self.validation_factory(residual_loader)
        if self.validation_factory is not None
        else self.options.validation
    )

    state.register_component(self.model_name, self.model)
    state.register_optimizer(self.optimizer_name, self.optimizer)
    state.parameter_roles[self.model_name] = ParameterRole.VARIANCE
    state.frozen_components = state.frozen_components | {self.mean_model_name}
    state.active_stage = self.name
    self._residual_loader = residual_loader
    self._validation = validation

train

train(state: TrainingState) -> StageResult

Train the variance model and transition to VARIANCE_READY.

With an early stopper and a checkpoint store configured, the stage then restores its best checkpoint, so the live variance model holds the best epoch's weights, while the mean model and, if registered, the mean optimizer stay registered. It saves the checkpoint again under the same key as a finalized VARIANCE_READY checkpoint with metadata {"stage": "variance", "stage_complete": True}.

Parameters:

Name	Type	Description	Default
state	TrainingState	Prepared state retaining a frozen mean component.	required

Returns:

Type	Description
StageResult	The supervised runner result, or the restored epoch's training
StageResult	metrics when the best checkpoint was restored.

Raises:

Type	Description
ValueError	If preparation is incomplete or training is non-finite.

Source code in src/stagereg/jax/supervised_staged.py

def train(self, state: TrainingState) -> StageResult:
    """Train the variance model and transition to ``VARIANCE_READY``.

    With an early stopper and a checkpoint store configured, the stage then
    restores its best checkpoint, so the live variance model holds the best
    epoch's weights, while the mean model and, if registered, the mean
    optimizer stay registered. It saves the checkpoint again under the same
    key as a finalized ``VARIANCE_READY`` checkpoint with metadata
    ``{"stage": "variance", "stage_complete": True}``.

    Args:
        state: Prepared state retaining a frozen mean component.

    Returns:
        The supervised runner result, or the restored epoch's training
        metrics when the best checkpoint was restored.

    Raises:
        ValueError: If preparation is incomplete or training is non-finite.
    """
    if state.lifecycle_state is not StageState.MEAN_READY:
        raise ValueError("variance stage requires MEAN_READY lifecycle state.")
    if self._residual_loader is None:
        raise ValueError("variance stage must be prepared before training.")
    return _train_stage(
        self,
        state,
        train_loader=self._residual_loader,
        validation=self._validation,
        ready=StageState.VARIANCE_READY,
    )

restore

restore(
    state: TrainingState, checkpoint: Checkpoint
) -> None

Restore the variance stage's finalized checkpoint into state.

Call it after the mean stage's restore. The live variance model and optimizer take the checkpoint's weights in place and are registered under model_name and optimizer_name. The mean model stays registered under mean_model_name, and the mean optimizer under mean_optimizer_name if it was registered before. Afterwards the state passes validate.

Parameters:

Name	Type	Description	Default
state	TrainingState	Training state the mean stage has been restored into.	required
checkpoint	Checkpoint	The variance stage's finalized checkpoint.	required

Raises:

Type	Description
ValueError	If checkpoint is not finalized by this stage, that is its lifecycle state is not VARIANCE_READY or its metadata lacks "stage": "variance" and "stage_complete": True, if it was saved with another model_name or mean_model_name, or if no mean model is registered under mean_model_name; all are checked before state or the live objects change. Also raised, as by restore_checkpoint, if the live model or optimizer is incompatible with the checkpoint, which may leave them and state partially restored.
TypeError	As by restore_checkpoint, if the checkpoint does not hold NNX snapshots or a JAX random key.

Source code in src/stagereg/jax/supervised_staged.py

def restore(self, state: TrainingState, checkpoint: Checkpoint) -> None:
    """Restore the variance stage's finalized checkpoint into ``state``.

    Call it after the mean stage's
    [`restore`][stagereg.jax.MeanStage.restore]. The live variance model and
    optimizer take the checkpoint's weights in place and are registered
    under ``model_name`` and ``optimizer_name``. The mean model stays
    registered under ``mean_model_name``, and the mean optimizer under
    ``mean_optimizer_name`` if it was registered before. Afterwards the
    state passes [`validate`][stagereg.jax.GammaVarianceStage.validate].

    Args:
        state: Training state the mean stage has been restored into.
        checkpoint: The variance stage's finalized checkpoint.

    Raises:
        ValueError: If ``checkpoint`` is not finalized by this stage, that
            is its lifecycle state is not ``VARIANCE_READY`` or its
            metadata lacks ``"stage": "variance"`` and ``"stage_complete":
            True``, if it was saved with another ``model_name`` or
            ``mean_model_name``, or if no mean model is registered under
            ``mean_model_name``; all are checked before ``state`` or the
            live objects change. Also raised, as by
            [`restore_checkpoint`][stagereg.jax.restore_checkpoint], if the
            live model or optimizer is incompatible with the checkpoint,
            which may leave them and ``state`` partially restored.
        TypeError: As by
            [`restore_checkpoint`][stagereg.jax.restore_checkpoint], if the
            checkpoint does not hold NNX snapshots or a JAX random key.
    """
    require_finalized(checkpoint, stage=self.name, ready=StageState.VARIANCE_READY)
    require_component_names(
        checkpoint,
        model_name=self.model_name,
        role=ParameterRole.VARIANCE,
        frozen=self.mean_model_name,
    )
    if self.mean_model_name not in state.model_components:
        raise ValueError(
            f"mean model component {self.mean_model_name!r} is not registered."
        )
    self._restore_live(state, checkpoint)

validate

validate(state: TrainingState) -> ValidationResult

Validate variance-stage lifecycle, ownership, and freezing.

Parameters:

Name	Type	Description	Default
state	TrainingState	Shared staged training state.	required

Returns:

Type	Description
ValidationResult	A validation result describing whether Step 2 is ready.

Source code in src/stagereg/jax/supervised_staged.py

def validate(self, state: TrainingState) -> ValidationResult:
    """Validate variance-stage lifecycle, ownership, and freezing.

    Args:
        state: Shared staged training state.

    Returns:
        A validation result describing whether Step 2 is ready.
    """
    passed = (
        state.lifecycle_state is StageState.VARIANCE_READY
        and state.model_components.get(self.model_name) is self.model
        and state.optimizer_states.get(self.optimizer_name) is self.optimizer
        and state.parameter_roles.get(self.model_name) is ParameterRole.VARIANCE
        and self.mean_model_name in state.frozen_components
    )
    return ValidationResult(
        passed=passed,
        message=None if passed else "variance stage invariants are not satisfied.",
    )

select_checkpoint

select_checkpoint(state: TrainingState) -> CheckpointRef

Return a reference to the variance stage's best checkpoint.

Parameters:

Name	Type	Description	Default
state	TrainingState	Shared staged training state.	required

Returns:

Type	Description
CheckpointRef	Reference to the checkpoint key the stage saves under: the
CheckpointRef	configured checkpoint_key, or f"{stage}/best" when none
CheckpointRef	was configured.

Raises:

Type	Description
ValueError	If no checkpoint exists under that key.

Source code in src/stagereg/jax/supervised_staged.py

def select_checkpoint(self, state: TrainingState) -> CheckpointRef:
    """Return a reference to the variance stage's best checkpoint.

    Args:
        state: Shared staged training state.

    Returns:
        Reference to the checkpoint key the stage saves under: the
        configured ``checkpoint_key``, or ``f"{stage}/best"`` when none
        was configured.

    Raises:
        ValueError: If no checkpoint exists under that key.
    """
    del state
    return _select_checkpoint(self)

stagereg.jax.validation

HeldOutValidation dataclass

HeldOutValidation(
    model: Module,
    loader: LoaderFactory,
    loss: SupervisedLoss,
    metrics: MetricSuite = MetricSuite(),
    metadata: dict[str, object] = dict(),
)

Evaluate an NNX model against a held-out validation loader.

Attributes:

Name	Type	Description
model	Module	The NNX module to evaluate.
loader	LoaderFactory	Factory producing the validation batch loader for a given split and epoch.
loss	SupervisedLoss	The supervised loss function used for evaluation.
metrics	MetricSuite	Registered batch/epoch metrics for validation.
metadata	dict[str, object]	Arbitrary caller-supplied metadata attached to this strategy.

__call__

__call__(
    state: TrainingState, *, epoch: int
) -> ValidationResult

Evaluate the held-out split after a training epoch.

Parameters:

Name	Type	Description	Default
state	TrainingState	The current training state, whose rng_state is advanced in place as validation batches are consumed.	required
epoch	int	The epoch just completed, passed to loader to select the validation batches for this epoch.	required

Returns:

Type	Description
ValidationResult	A ValidationResult with bare
ValidationResult	metric names and passed=True.

Raises:

Type	Description
TypeError	If state.rng_state is not a JAX random key.

Source code in src/stagereg/jax/validation.py

def __call__(self, state: TrainingState, *, epoch: int) -> ValidationResult:
    """Evaluate the held-out split after a training epoch.

    Args:
        state: The current training state, whose ``rng_state`` is
            advanced in place as validation batches are consumed.
        epoch: The epoch just completed, passed to ``loader`` to
            select the validation batches for this epoch.

    Returns:
        A [`ValidationResult`][stagereg.core.ValidationResult] with bare
        metric names and ``passed=True``.

    Raises:
        TypeError: If ``state.rng_state`` is not a JAX random key.
    """
    if not isinstance(state.rng_state, jax.Array):
        raise TypeError("state.rng_state must be a JAX random key.")
    metrics, state.rng_state = evaluate_loader(
        self.model,
        self.loader(split="validation", epoch=epoch),
        key=state.rng_state,
        metrics=self.metrics,
        loss=self.loss,
    )
    return ValidationResult(passed=True, metrics=metrics, message=None)