stagereg

stagereg is a library for stage-oriented probabilistic regression: training models that predict a distribution over the target, not just a point estimate. Its core is a set of backend-neutral contracts for batches, predictive distributions, losses, metrics, checkpoints, early stopping and tracking. An optional JAX backend built on Flax NNX and Optax implements them.

Why train in stages

Splitting training into separate stages fixes two problems that fitting everything at once leaves open:

• It avoids the pathologies of joint training. Mean-variance estimation (Nix & Weigend, 1994) fits a mean and a variance together under one Gaussian negative log-likelihood. That couples the two fits: the variance can grow to excuse a poor mean fit, and the mean stalls where the variance is large (Detlefsen et al., 2019). Fitting the mean before the variance removes the coupling; warm-up schemes that fit the mean first (Sluijterman et al., 2024) apply the same idea for part of training.
• It disentangles aleatoric and epistemic uncertainty. Once its own stage has fixed the aleatoric variance, a later posterior stage can attribute the remaining spread of mean functions to epistemic variance, so neither absorbs the other (Yi & Bessa, 2025).

How stagereg trains in stages

A run is a sequence of named stages: a mean stage, a variance stage, and an optional posterior stage. Each stage checks that the lifecycle of the shared training state allows it to run, trains and validates, and hands on a finalized checkpoint the next stage resumes from. Every metric a run reports is tagged with the stage that produced it. See two-step training for the mean and variance stages and the posterior stage for the third.

Install

Install with the JAX backend (JAX, Flax and Optax):

pip install "stagereg[jax]"

The optional plot extra also installs matplotlib, for plotting results:

pip install "stagereg[jax,plot]"

With uv, use the equivalent uv add form:

uv add "stagereg[jax]"
uv add "stagereg[jax,plot]"

Drop the [jax] extra for the backend-neutral contracts alone, which depend only on NumPy.

The documentation tracks the main branch, with no per-release docs before 1.0, so it may describe features newer than your installed version.

Minimal example

Fit a linear model to y = 3x + 2 for a fixed number of epochs with the JAX backend:

import jax
import jax.numpy as jnp
import optax
from flax import nnx

from stagereg.core.types import Batch
from stagereg.jax import create_optimizer, initialize_training_state, run_supervised


def squared_error(model, inputs, targets, sample_weight, key, training):
    return jnp.mean(jnp.square(model(inputs) - targets))


inputs = jnp.linspace(-1.0, 1.0, 32).reshape(-1, 1)
targets = 3.0 * inputs + 2.0


def loader(*, split, epoch):
    return [Batch(inputs=inputs, targets=targets)]


model = nnx.Linear(1, 1, rngs=nnx.Rngs(0))
optimizer = create_optimizer(model, optax.sgd(learning_rate=0.1))
state = initialize_training_state(model, optimizer, rng_key=jax.random.key(0))

result = run_supervised(
    model=model,
    optimizer=optimizer,
    train_loader=loader,
    loss=squared_error,
    state=state,
    epochs=3,
)
print(f"final training loss: {result.loss:.4f}")

Contributing

Development setup, the checks a change must pass and how to build this site are in CONTRIBUTING.md.