`stagereg.core`¶

Every symbol exported from stagereg.core, the backend-neutral contracts and utilities, grouped by the module that defines it.

`stagereg.core.checkpoints`¶

Checkpoint `dataclass` ¶

Checkpoint(
    state: TrainingState,
    epoch: int,
    parameters: Any | None = None,
    optimizer_state: Any | None = None,
    rng_state: Any | None = None,
    early_stopping_state: Any | None = None,
    metadata: Mapping[str, Any] = dict(),
)

The complete state necessary to resume staged training.

Attributes:

Name	Type	Description
`state`	`TrainingState`	Snapshot of the `TrainingState` at the time the checkpoint was saved.
`epoch`	`int`	Epoch of the active stage at which the checkpoint was saved.
`parameters`	`Any \| None`	Snapshot of the model's parameters, in the backend's own format. Defaults to `None`.
`optimizer_state`	`Any \| None`	Snapshot of the optimizer's state, in the backend's own format. Defaults to `None`.
`rng_state`	`Any \| None`	Random key to resume training with. It mirrors the `rng_state` of the embedded `state`. Defaults to `None`.
`early_stopping_state`	`Any \| None`	State of the early stopper, such as an `EarlyStoppingState`, so that patience and the best value carry over on resume. Defaults to `None`.
`metadata`	`Mapping[str, Any]`	Information about the checkpoint, such as the stage that saved it. Defaults to an empty mapping.

CheckpointStore ¶

Bases: Protocol

Persists checkpoints under opaque keys.

InMemoryCheckpointStore ¶

InMemoryCheckpointStore()

A CheckpointStore kept in process memory.

Checkpoints live only for the lifetime of the store, which makes it the natural choice for examples, tests, and short runs that only need the best checkpoint of the current process restored at the end. Saving under an existing key replaces the checkpoint held there.

Source code in src/stagereg/core/checkpoints.py

def __init__(self) -> None:
    self._checkpoints: dict[str, Checkpoint] = {}

save ¶

save(key: str, checkpoint: Checkpoint) -> None

Store a checkpoint under a key, replacing any checkpoint there.

Parameters:

Name	Type	Description	Default
`key`	`str`	The opaque key to store the checkpoint under.	required
`checkpoint`	`Checkpoint`	The checkpoint to store.	required

Returns:

Type	Description
`None`	None.

Source code in src/stagereg/core/checkpoints.py

def save(self, key: str, checkpoint: Checkpoint) -> None:
    """Store a checkpoint under a key, replacing any checkpoint there.

    Args:
        key: The opaque key to store the checkpoint under.
        checkpoint: The checkpoint to store.

    Returns:
        None.
    """
    self._checkpoints[key] = checkpoint

load ¶

load(key: str) -> Checkpoint

Return the checkpoint stored under a key.

Parameters:

Name	Type	Description	Default
`key`	`str`	The key the checkpoint was saved under.	required

Returns:

Type	Description
`Checkpoint`	The stored checkpoint.

Raises:

Type	Description
`KeyError`	If no checkpoint has been saved under `key`.

Source code in src/stagereg/core/checkpoints.py

def load(self, key: str) -> Checkpoint:
    """Return the checkpoint stored under a key.

    Args:
        key: The key the checkpoint was saved under.

    Returns:
        The stored checkpoint.

    Raises:
        KeyError: If no checkpoint has been saved under ``key``.
    """
    return self._checkpoints[key]

exists ¶

exists(key: str) -> bool

Report whether a checkpoint is stored under a key.

Parameters:

Name	Type	Description	Default
`key`	`str`	The key to look up.	required

Returns:

Type	Description
`bool`	`True` if a checkpoint is stored under `key`.

Source code in src/stagereg/core/checkpoints.py

def exists(self, key: str) -> bool:
    """Report whether a checkpoint is stored under a key.

    Args:
        key: The key to look up.

    Returns:
        ``True`` if a checkpoint is stored under ``key``.
    """
    return key in self._checkpoints

`stagereg.core.distributions`¶

PredictiveDistribution ¶

Bases: Protocol

A probability distribution produced by a regression model.

This protocol is the backend-neutral model/loss boundary. Epoch metrics use backend adapters to materialize host arrays rather than retaining distribution objects, which intentionally have no concatenation or indexing contract.

DistributionHead ¶

Bases: Protocol

Produces a predictive distribution from model features.

Likelihood ¶

Bases: Protocol

Evaluates a target under a predictive distribution.

DistributionLoss ¶

Bases: Protocol

Computes an unreduced loss from a predictive distribution.

PredictionLoss ¶

Bases: Protocol

Computes an unreduced loss from deterministic array predictions.

Loss `module-attribute` ¶

Loss = DistributionLoss

`stagereg.core.early_stopping`¶

OptimizationMode ¶

Bases: StrEnum

The direction in which a monitored metric is considered improved.

EarlyStoppingState `dataclass` ¶

EarlyStoppingState(
    metric: str,
    mode: OptimizationMode,
    min_delta: float,
    patience: int,
    source: Split = Split.VALIDATION,
    best_score: float | None = None,
    best_epoch: int | None = None,
    non_improving_epochs: int = 0,
    stopped: bool = False,
)

Checkpointable state for a configured early-stopping policy.

Attributes:

Name	Type	Description
`metric`	`str`	Name of the monitored metric.
`mode`	`OptimizationMode`	Direction in which the metric is considered improved.
`min_delta`	`float`	Minimum change required to qualify as an improvement.
`patience`	`int`	Number of consecutive non-improving epochs tolerated before stopping.
`source`	`Split`	Split the monitored metric is measured on.
`best_score`	`float \| None`	Best value observed so far, or `None` if no value has been observed yet.
`best_epoch`	`int \| None`	Epoch at which `best_score` was observed, or `None` if no value has been observed yet.
`non_improving_epochs`	`int`	Number of consecutive epochs without an improvement.
`stopped`	`bool`	Whether patience has been exhausted and training should stop.

EarlyStoppingDecision `dataclass` ¶

EarlyStoppingDecision(
    state: EarlyStoppingState, improved: bool
)

The result of observing one monitored metric value.

Attributes:

Name	Type	Description
`state`	`EarlyStoppingState`	Updated early-stopping state after the observation.
`improved`	`bool`	Whether the observed value improved on the best score.

should_stop `property` ¶

should_stop: bool

Whether the configured patience has been exhausted.

EarlyStopper ¶

EarlyStopper(
    *,
    metric: str,
    mode: OptimizationMode | str,
    patience: int,
    min_delta: float = 0.0,
    source: Split | str = Split.VALIDATION,
    state: EarlyStoppingState | None = None,
)

Tracks a scalar metric and decides when training should stop.

Initialize an early stopper.

Parameters:

Name	Type	Description	Default
`metric`	`str`	Name of the metric to monitor.	required
`mode`	`OptimizationMode \| str`	Direction in which the metric is considered improved, either `"min"` or `"max"`.	required
`patience`	`int`	Number of consecutive non-improving epochs tolerated before stopping.	required
`min_delta`	`float`	Minimum change required to qualify as an improvement. Defaults to `0.0`.	`0.0`
`source`	`Split \| str`	Split the monitored metric is measured on, either `"train"` or `"validation"`. Defaults to `Split.VALIDATION`.	`VALIDATION`
`state`	`EarlyStoppingState \| None`	Previously checkpointed state to resume from. Must match the configuration passed to this constructor.	`None`

Raises:

Type	Description
`ValueError`	If `metric` is empty, `patience` is negative, `min_delta` is not a finite non-negative number, `mode` or `source` is not a recognised value, or `state` does not match the given configuration.

Source code in src/stagereg/core/early_stopping.py

def __init__(
    self,
    *,
    metric: str,
    mode: OptimizationMode | str,
    patience: int,
    min_delta: float = 0.0,
    source: Split | str = Split.VALIDATION,
    state: EarlyStoppingState | None = None,
) -> None:
    """Initialize an early stopper.

    Args:
        metric: Name of the metric to monitor.
        mode: Direction in which the metric is considered improved,
            either ``"min"`` or ``"max"``.
        patience: Number of consecutive non-improving epochs
            tolerated before stopping.
        min_delta: Minimum change required to qualify as an
            improvement. Defaults to ``0.0``.
        source: Split the monitored metric is measured on, either
            ``"train"`` or ``"validation"``. Defaults to
            ``Split.VALIDATION``.
        state: Previously checkpointed state to resume from. Must
            match the configuration passed to this constructor.

    Raises:
        ValueError: If ``metric`` is empty, ``patience`` is
            negative, ``min_delta`` is not a finite non-negative
            number, ``mode`` or ``source`` is not a recognised value,
            or ``state`` does not match the given configuration.
    """
    if not metric:
        raise ValueError("metric must be a non-empty string.")
    if patience < 0:
        raise ValueError("patience must be non-negative.")
    if not isfinite(min_delta) or min_delta < 0:
        raise ValueError("min_delta must be finite and non-negative.")

    normalized_mode = OptimizationMode(mode)
    normalized_source = Split(source)
    self._state = state or EarlyStoppingState(
        metric=metric,
        mode=normalized_mode,
        min_delta=min_delta,
        patience=patience,
        source=normalized_source,
    )
    if (
        self._state.metric != metric
        or self._state.mode != normalized_mode
        or self._state.min_delta != min_delta
        or self._state.patience != patience
        or self._state.source != normalized_source
    ):
        raise ValueError("state does not match the early-stopper configuration.")

state `property` ¶

state: EarlyStoppingState

Return the current checkpointable stopping state.

expects_validation ¶

expects_validation() -> bool

Whether the monitored metric must come from validation.

Returns:

Type	Description
`bool`	`True` if the metric configured for this stopper is
`bool`	measured on the validation split (i.e. requires a
`bool`	`ValidationStrategy` to be
`bool`	supplied to the training loop), `False` if it is measured
`bool`	on the training split.

Source code in src/stagereg/core/early_stopping.py

def expects_validation(self) -> bool:
    """Whether the monitored metric must come from validation.

    Returns:
        ``True`` if the metric configured for this stopper is
        measured on the validation split (i.e. requires a
        [`ValidationStrategy`][stagereg.core.ValidationStrategy] to be
        supplied to the training loop), ``False`` if it is measured
        on the training split.
    """
    return self._state.source is Split.VALIDATION

monitored_metric_name ¶

monitored_metric_name() -> str

Return the name of the metric this stopper monitors.

Returns:

Type	Description
`str`	The metric name passed to the constructor, used by callers
`str`	to look up the corresponding value in per-epoch metrics.

Source code in src/stagereg/core/early_stopping.py

def monitored_metric_name(self) -> str:
    """Return the name of the metric this stopper monitors.

    Returns:
        The metric name passed to the constructor, used by callers
        to look up the corresponding value in per-epoch metrics.
    """
    return self._state.metric

observe ¶

observe(
    value: float, *, epoch: int
) -> EarlyStoppingDecision

Record a metric observed after epoch and return its decision.

Parameters:

Name	Type	Description	Default
`value`	`float`	The monitored metric value observed for this epoch.	required
`epoch`	`int`	The epoch at which `value` was observed.	required

Returns:

Type	Description
`EarlyStoppingDecision`	An `EarlyStoppingDecision`
`EarlyStoppingDecision`	with the updated state and whether `value` improved on the best
`EarlyStoppingDecision`	score so far.

Raises:

Type	Description
`ValueError`	If `epoch` is negative or `value` is not finite.
`RuntimeError`	If early stopping has already been triggered.

Source code in src/stagereg/core/early_stopping.py

def observe(self, value: float, *, epoch: int) -> EarlyStoppingDecision:
    """Record a metric observed after ``epoch`` and return its decision.

    Args:
        value: The monitored metric value observed for this epoch.
        epoch: The epoch at which ``value`` was observed.

    Returns:
        An [`EarlyStoppingDecision`][stagereg.core.EarlyStoppingDecision]
        with the updated state and whether ``value`` improved on the best
        score so far.

    Raises:
        ValueError: If ``epoch`` is negative or ``value`` is not
            finite.
        RuntimeError: If early stopping has already been triggered.
    """
    if epoch < 0:
        raise ValueError("epoch must be non-negative.")
    if not isfinite(value):
        raise ValueError("observed metric value must be finite.")
    if self._state.stopped:
        raise RuntimeError("cannot observe a metric after early stopping.")

    improved = self._is_improvement(value)
    if improved:
        self._state = EarlyStoppingState(
            metric=self._state.metric,
            mode=self._state.mode,
            min_delta=self._state.min_delta,
            patience=self._state.patience,
            source=self._state.source,
            best_score=value,
            best_epoch=epoch,
        )
    else:
        non_improving_epochs = self._state.non_improving_epochs + 1
        self._state = EarlyStoppingState(
            metric=self._state.metric,
            mode=self._state.mode,
            min_delta=self._state.min_delta,
            patience=self._state.patience,
            source=self._state.source,
            best_score=self._state.best_score,
            best_epoch=self._state.best_epoch,
            non_improving_epochs=non_improving_epochs,
            stopped=non_improving_epochs > self._state.patience,
        )
    return EarlyStoppingDecision(state=self._state, improved=improved)

`stagereg.core.losses`¶

NegativeLogLikelihoodLoss `dataclass` ¶

NegativeLogLikelihoodLoss(
    beta: float = 0.0,
    stop_gradient: Callable[[Array], Array] = _identity,
    target_transform: Callable[[Array], Array] = _identity,
)

Configurable negative log-likelihood for predictive distributions.

Attributes:

Name	Type	Description
`beta`	`float`	Variance-reweighting exponent in `[0, 1]`. Zero recovers ordinary negative log-likelihood.
`stop_gradient`	`Callable[[Array], Array]`	Callable excluding the variance weight from gradient propagation. Defaults to the identity, which lets gradients flow through the weight. With JAX and `beta > 0`, pass `jax.lax.stop_gradient`; otherwise a different objective from beta-NLL is trained.
`target_transform`	`Callable[[Array], Array]`	Callable applied to targets before evaluating the predictive distribution. Defaults to the identity transform.

per_example ¶

per_example(
    prediction: PredictiveDistribution, batch: Batch
) -> Array

Compute the optionally transformed and beta-weighted NLL.

Parameters:

Name	Type	Description	Default
`prediction`	`PredictiveDistribution`	The predictive distribution produced for `batch`.	required
`batch`	`Batch`	Batch whose targets are scored under `prediction`.	required

Returns:

Type	Description
`Array`	An array of per-example negative log-likelihood values.

Raises:

Type	Description
`ValueError`	If `batch.targets` is missing.

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

def per_example(self, prediction: PredictiveDistribution, batch: Batch) -> Array:
    """Compute the optionally transformed and beta-weighted NLL.

    Args:
        prediction: The predictive distribution produced for ``batch``.
        batch: Batch whose targets are scored under ``prediction``.

    Returns:
        An array of per-example negative log-likelihood values.

    Raises:
        ValueError: If ``batch.targets`` is missing.
    """
    if batch.targets is None:
        raise ValueError("batch.targets must be provided.")
    targets = self.target_transform(batch.targets)
    nll = -prediction.log_prob(targets)
    if self.beta == 0.0:
        return nll
    weight = self.stop_gradient(prediction.variance()) ** self.beta
    return nll * weight

SquaredErrorLoss `dataclass` ¶

SquaredErrorLoss()

Per-example squared error for deterministic predictions.

per_example ¶

per_example(prediction: Array, batch: Batch) -> Array

Compute squared prediction errors.

Parameters:

Name	Type	Description	Default
`prediction`	`Array`	Deterministic model predictions.	required
`batch`	`Batch`	Batch containing regression targets.	required

Returns:

Type	Description
`Array`	Elementwise squared prediction errors.

Raises:

Type	Description
`ValueError`	If `batch.targets` is missing.

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

def per_example(self, prediction: Array, batch: Batch) -> Array:
    """Compute squared prediction errors.

    Args:
        prediction: Deterministic model predictions.
        batch: Batch containing regression targets.

    Returns:
        Elementwise squared prediction errors.

    Raises:
        ValueError: If ``batch.targets`` is missing.
    """
    if batch.targets is None:
        raise ValueError("batch.targets must be provided.")
    if prediction.shape != batch.targets.shape:
        raise ValueError("prediction and batch.targets must have matching shapes.")
    return (prediction - batch.targets) ** 2

add_epsilon ¶

add_epsilon(
    epsilon: float = 1e-12,
) -> Callable[[Array], Array]

Create a target transform that adds a positive stabilization offset.

Parameters:

Name	Type	Description	Default
`epsilon`	`float`	Positive finite offset added to each target.	`1e-12`

Returns:

Type	Description
`Callable[[Array], Array]`	A callable adding `epsilon` to an array.

Raises:

Type	Description
`ValueError`	If `epsilon` is not positive and finite.

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

def add_epsilon(epsilon: float = 1e-12) -> Callable[[Array], Array]:
    """Create a target transform that adds a positive stabilization offset.

    Args:
        epsilon: Positive finite offset added to each target.

    Returns:
        A callable adding ``epsilon`` to an array.

    Raises:
        ValueError: If ``epsilon`` is not positive and finite.
    """
    if not math.isfinite(epsilon) or epsilon <= 0.0:
        raise ValueError("epsilon must be positive and finite.")

    def transform(targets: Array) -> Array:
        return targets + epsilon

    return transform

GaussianNLLLoss `module-attribute` ¶

GaussianNLLLoss = NegativeLogLikelihoodLoss

Alias of NegativeLogLikelihoodLoss.

Kept for compatibility with the pre-consolidation name; prefer the canonical name.

BetaNLLLoss `module-attribute` ¶

BetaNLLLoss = NegativeLogLikelihoodLoss

Alias of NegativeLogLikelihoodLoss.

Kept for compatibility with the pre-consolidation name; prefer the canonical name. The beta-NLL weighting is the beta argument of NegativeLogLikelihoodLoss.

`stagereg.core.metric_registry`¶

EvaluationGrid `dataclass` ¶

EvaluationGrid(values: FloatArray)

A shared scalar CRPS integration grid.

Attributes:

Name	Type	Description
`values`	`FloatArray`	Finite, strictly increasing integration coordinates.

PredictionInterval `dataclass` ¶

PredictionInterval(
    level: float, lower: FloatArray, upper: FloatArray
)

Host-resident bounds labelled by their central confidence level.

Attributes:

Name	Type	Description
`level`	`float`	Central confidence level of the interval, strictly between 0 and 1.
`lower`	`FloatArray`	Lower bounds, one per scoring unit.
`upper`	`FloatArray`	Upper bounds, one per scoring unit, each no smaller than its lower bound.

EpochPredictionData `dataclass` ¶

EpochPredictionData(
    targets: FloatArray,
    mean: FloatArray,
    variance: FloatArray | None = None,
    predictive_samples: FloatArray | None = None,
    reference_samples: FloatArray | None = None,
    intervals: tuple[PredictionInterval, ...] = (),
    coordinate: FloatArray | None = None,
    evaluation_grid: EvaluationGrid | None = None,
    log_likelihood: FloatArray | None = None,
)

Typed, host-materialized scalar predictions for one batch or epoch.

Rows are independent scalar scoring units. Sample matrices use rows for scoring units and columns for samples; non-scalar distribution events are intentionally not represented by this contract.

Attributes:

Name	Type	Description
`targets`	`FloatArray`	Observed target, one per scoring unit.
`mean`	`FloatArray`	Predictive mean, one per scoring unit.
`variance`	`FloatArray \| None`	Non-negative predictive variance, one per scoring unit. Defaults to `None` when no metric requires it.
`predictive_samples`	`FloatArray \| None`	Samples from the predictive distribution, one row per scoring unit and one column per sample. Defaults to `None` when no metric requires them.
`reference_samples`	`FloatArray \| None`	Samples from each scoring unit's reference distribution, the empirical distribution the predictive samples are scored against, one row per scoring unit and one column per sample. Defaults to `None` when no metric requires them.
`intervals`	`tuple[PredictionInterval, ...]`	Prediction intervals, each a `PredictionInterval` at a distinct confidence level. Defaults to no intervals.
`coordinate`	`FloatArray \| None`	Numeric coordinate of each scoring unit, such as time or remaining useful life, that a metric orders or weights scoring units by. Defaults to `None` when no metric requires it.
`evaluation_grid`	`EvaluationGrid \| None`	Shared `EvaluationGrid` that distributional scores are integrated over. Defaults to `None` when no metric requires it.
`log_likelihood`	`FloatArray \| None`	Log-density of each scoring unit's target under the predictive distribution. Defaults to `None` when no metric requires it.

interval ¶

interval(level: float) -> PredictionInterval

Return the prediction interval registered at exactly level.

Parameters:

Name	Type	Description	Default
`level`	`float`	Required central confidence level.	required

Returns:

Type	Description
`PredictionInterval`	The matching labelled interval.

Raises:

Type	Description
`ValueError`	If no interval at `level` is present.

Source code in src/stagereg/core/metric_registry.py

def interval(self, level: float) -> PredictionInterval:
    """Return the prediction interval registered at exactly ``level``.

    Args:
        level: Required central confidence level.

    Returns:
        The matching labelled interval.

    Raises:
        ValueError: If no interval at ``level`` is present.
    """
    for interval in self.intervals:
        if interval.level == level:
            return interval
    raise ValueError(f"prediction interval at level {level} is required.")

MetricRequirements `dataclass` ¶

MetricRequirements(
    variance: bool = False,
    predictive_samples: bool = False,
    reference_samples: bool = False,
    interval_levels: frozenset[float] = frozenset(),
    coordinate: bool = False,
    evaluation_grid: bool = False,
    log_likelihood: bool = False,
)

Materialized fields required by one or more epoch metrics.

Targets and predictive means are always materialized; each attribute asks for one more field of EpochPredictionData.

Attributes:

Name	Type	Description
`variance`	`bool`	Whether the predictive variance is required. Defaults to `False`.
`predictive_samples`	`bool`	Whether samples from the predictive distribution are required. Defaults to `False`.
`reference_samples`	`bool`	Whether samples from each scoring unit's reference distribution are required. Defaults to `False`.
`interval_levels`	`frozenset[float]`	Confidence levels, each in `(0, 1)`, at which prediction intervals are required. Defaults to no levels.
`coordinate`	`bool`	Whether each scoring unit's numeric coordinate is required. Defaults to `False`.
`evaluation_grid`	`bool`	Whether the shared evaluation grid is required. Defaults to `False`.
`log_likelihood`	`bool`	Whether each target's predictive log-density is required. Defaults to `False`.

union ¶

union(other: MetricRequirements) -> MetricRequirements

Combine this requirement set with another.

Parameters:

Name	Type	Description	Default
`other`	`MetricRequirements`	Additional requirements.	required

Returns:

Type	Description
`MetricRequirements`	A requirement set containing the union of both inputs.

Source code in src/stagereg/core/metric_registry.py

def union(self, other: MetricRequirements) -> MetricRequirements:
    """Combine this requirement set with another.

    Args:
        other: Additional requirements.

    Returns:
        A requirement set containing the union of both inputs.
    """
    return MetricRequirements(
        variance=self.variance or other.variance,
        predictive_samples=self.predictive_samples or other.predictive_samples,
        reference_samples=self.reference_samples or other.reference_samples,
        interval_levels=self.interval_levels | other.interval_levels,
        coordinate=self.coordinate or other.coordinate,
        evaluation_grid=self.evaluation_grid or other.evaluation_grid,
        log_likelihood=self.log_likelihood or other.log_likelihood,
    )

EpochMetric ¶

Bases: Protocol

A host-side scalar metric over typed epoch prediction data.

name `property` ¶

name: str

Return the emitted metric name.

requirements `property` ¶

requirements: MetricRequirements

Return fields that must be materialized for this metric.

call ¶

__call__(data: EpochPredictionData) -> float

Compute the scalar metric value.

Source code in src/stagereg/core/metric_registry.py

def __call__(self, data: EpochPredictionData, /) -> float:
    """Compute the scalar metric value."""
    ...

RootMeanSquaredError `dataclass` ¶

RootMeanSquaredError(name: str = 'rmse')

Epoch adapter for root mean squared error.

Attributes:

Name	Type	Description
`name`	`str`	Bare metric name the value is reported under. Defaults to `rmse`.

requirements `property` ¶

requirements: MetricRequirements

Return the base target/mean requirements.

call ¶

__call__(data: EpochPredictionData) -> float

Compute RMSE from scalar targets and predictive means.

Source code in src/stagereg/core/metric_registry.py

def __call__(self, data: EpochPredictionData, /) -> float:
    """Compute RMSE from scalar targets and predictive means."""
    return rmse(data.targets, data.mean)

NegativeLogLikelihood `dataclass` ¶

NegativeLogLikelihood(name: str = 'nll')

Mean negative log-density of the targets under the predictive distribution.

The predictor computes each target's exact log-density, so a non-Gaussian predictive (such as a mixture over draws) is scored without a Gaussian approximation.

Attributes:

Name	Type	Description
`name`	`str`	Bare metric name the value is reported under. Defaults to `nll`.

requirements `property` ¶

requirements: MetricRequirements

Require each target's predictive log-density.

call ¶

__call__(data: EpochPredictionData) -> float

Average the negative log-likelihood over scoring units.

Raises:

Type	Description
`ValueError`	If the log-likelihood is absent.

Source code in src/stagereg/core/metric_registry.py

def __call__(self, data: EpochPredictionData, /) -> float:
    """Average the negative log-likelihood over scoring units.

    Raises:
        ValueError: If the log-likelihood is absent.
    """
    if data.log_likelihood is None:
        raise ValueError("NLL requires log_likelihood.")
    return float(-np.mean(data.log_likelihood))

IntervalCoverage `dataclass` ¶

IntervalCoverage(
    level: float = 0.95, name: str = "coverage"
)

Observed coverage of an explicitly labelled prediction interval.

Attributes:

Name	Type	Description
`level`	`float`	Confidence level of the labelled interval, in `(0, 1)`.
`name`	`str`	Bare metric name the value is reported under. Defaults to `coverage`.

requirements `property` ¶

requirements: MetricRequirements

Require an interval at the configured level.

call ¶

__call__(data: EpochPredictionData) -> float

Compute inclusive interval coverage.

Source code in src/stagereg/core/metric_registry.py

def __call__(self, data: EpochPredictionData, /) -> float:
    """Compute inclusive interval coverage."""
    interval = data.interval(self.level)
    return coverage(data.targets, interval.lower, interval.upper)

WeightedSpread `dataclass` ¶

WeightedSpread(level: float = 0.95, name: str = 'wsu')

Weighted spread for an interval and explicit numeric coordinate.

Attributes:

Name	Type	Description
`level`	`float`	Confidence level of the labelled interval, in `(0, 1)`.
`name`	`str`	Bare metric name the value is reported under. Defaults to `wsu`.

requirements `property` ¶

requirements: MetricRequirements

Require labelled bounds and a coordinate.

call ¶

__call__(data: EpochPredictionData) -> float

Compute normalized weighted interval spread.

Raises:

Type	Description
`ValueError`	If the required coordinate is absent.

Source code in src/stagereg/core/metric_registry.py

def __call__(self, data: EpochPredictionData, /) -> float:
    """Compute normalized weighted interval spread.

    Raises:
        ValueError: If the required coordinate is absent.
    """
    if data.coordinate is None:
        raise ValueError("WeightedSpread requires an explicit coordinate.")
    interval = data.interval(self.level)
    order = np.argsort(data.coordinate)
    return wsu(
        interval.lower[order],
        interval.upper[order],
        data.coordinate[order],
    )

PointContinuousRankedProbabilityScore `dataclass` ¶

PointContinuousRankedProbabilityScore(
    name: str = "point_crps",
)

Mean point-CRPS across independent scalar scoring units.

Attributes:

Name	Type	Description
`name`	`str`	Bare metric name the value is reported under. Defaults to `point_crps`.

requirements `property` ¶

requirements: MetricRequirements

Require predictive samples and the shared evaluation grid.

call ¶

__call__(data: EpochPredictionData) -> float

Compute one point-CRPS per unit, then average units equally.

Raises:

Type	Description
`ValueError`	If predictive samples or the grid are absent.

Source code in src/stagereg/core/metric_registry.py

def __call__(self, data: EpochPredictionData, /) -> float:
    """Compute one point-CRPS per unit, then average units equally.

    Raises:
        ValueError: If predictive samples or the grid are absent.
    """
    if data.predictive_samples is None:
        raise ValueError("Point CRPS requires predictive_samples.")
    if data.evaluation_grid is None:
        raise ValueError("Point CRPS requires an evaluation_grid.")
    scores = [
        point_crps(target, samples, data.evaluation_grid.values)
        for target, samples in zip(
            data.targets, data.predictive_samples, strict=True
        )
    ]
    return float(np.mean(scores))

SampleContinuousRankedProbabilityScore `dataclass` ¶

SampleContinuousRankedProbabilityScore(
    name: str = "sample_crps",
)

Mean exact CRPS of predictive samples across scalar scoring units.

Each unit's predictive samples are scored against its target with sample_crps, the energy form of the CRPS, so unlike PointContinuousRankedProbabilityScore it needs no evaluation grid.

Attributes:

Name	Type	Description
`name`	`str`	Bare metric name the value is reported under. Defaults to `sample_crps`.

requirements `property` ¶

requirements: MetricRequirements

Require predictive samples only.

call ¶

__call__(data: EpochPredictionData) -> float

Compute one CRPS per unit, then average units equally.

Raises:

Type	Description
`ValueError`	If predictive samples are absent.

Source code in src/stagereg/core/metric_registry.py

def __call__(self, data: EpochPredictionData, /) -> float:
    """Compute one CRPS per unit, then average units equally.

    Raises:
        ValueError: If predictive samples are absent.
    """
    if data.predictive_samples is None:
        raise ValueError("Sample CRPS requires predictive_samples.")
    scores = [
        sample_crps(target, samples)
        for target, samples in zip(
            data.targets, data.predictive_samples, strict=True
        )
    ]
    return float(np.mean(scores))

ContinuousRankedProbabilityScore `dataclass` ¶

ContinuousRankedProbabilityScore(name: str = 'crps')

Mean expected CRPS under per-unit empirical reference distributions.

Attributes:

Name	Type	Description
`name`	`str`	Bare metric name the value is reported under. Defaults to `crps`.

requirements `property` ¶

requirements: MetricRequirements

Require predictive/reference samples and the shared grid.

call ¶

__call__(data: EpochPredictionData) -> float

Compute expected empirical-reference CRPS and average scoring units.

Raises:

Type	Description
`ValueError`	If either sample matrix or the grid is absent.

Source code in src/stagereg/core/metric_registry.py

def __call__(self, data: EpochPredictionData, /) -> float:
    """Compute expected empirical-reference CRPS and average scoring units.

    Raises:
        ValueError: If either sample matrix or the grid is absent.
    """
    if data.predictive_samples is None:
        raise ValueError("CRPS requires predictive_samples.")
    if data.reference_samples is None:
        raise ValueError("CRPS requires reference_samples.")
    if data.evaluation_grid is None:
        raise ValueError("CRPS requires an evaluation_grid.")
    scores = [
        crps(reference, predictive, data.evaluation_grid.values)
        for reference, predictive in zip(
            data.reference_samples, data.predictive_samples, strict=True
        )
    ]
    return float(np.mean(scores))

`stagereg.core.metrics`¶

cdf ¶

cdf(x: float, samples: ArrayLike) -> float

Return the empirical cumulative distribution at x.

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

def cdf(x: float, samples: ArrayLike) -> float:
    """Return the empirical cumulative distribution at ``x``."""
    sample_vector = np.asarray(samples, dtype=float)
    if sample_vector.ndim != 1:
        raise ValueError("samples must be one-dimensional.")
    if sample_vector.size == 0:
        return 0.0
    if not np.isfinite(x) or not np.all(np.isfinite(sample_vector)):
        raise ValueError("x and samples must contain only finite values.")
    return float(np.mean(sample_vector <= x))

crps ¶

crps(
    samples_true: ArrayLike,
    samples_pred: ArrayLike,
    x_grid: ArrayLike,
) -> float

Approximate expected CRPS under a scalar empirical reference distribution.

Parameters:

Name	Type	Description	Default
`samples_true`	`ArrayLike`	Samples from the reference scalar distribution.	required
`samples_pred`	`ArrayLike`	Samples from the predictive scalar distribution.	required
`x_grid`	`ArrayLike`	Finite, strictly increasing integration coordinates.	required

Returns:

Type	Description
`float`	The predictive CRPS averaged over the empirical reference distribution.

Raises:

Type	Description
`ValueError`	If samples or the grid are empty, non-finite, or malformed.

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

def crps(samples_true: ArrayLike, samples_pred: ArrayLike, x_grid: ArrayLike) -> float:
    """Approximate expected CRPS under a scalar empirical reference distribution.

    Args:
        samples_true: Samples from the reference scalar distribution.
        samples_pred: Samples from the predictive scalar distribution.
        x_grid: Finite, strictly increasing integration coordinates.

    Returns:
        The predictive CRPS averaged over the empirical reference distribution.

    Raises:
        ValueError: If samples or the grid are empty, non-finite, or malformed.
    """
    grid = _evaluation_grid(x_grid)
    true_samples = _as_nonempty_vector(samples_true, name="samples_true")
    predicted_samples = _as_nonempty_vector(samples_pred, name="samples_pred")
    true_cdf = np.mean(true_samples[None, :] <= grid[:, None], axis=1)
    predicted_cdf = np.mean(predicted_samples[None, :] <= grid[:, None], axis=1)
    expected_squared_distance = (
        true_cdf * (1.0 - predicted_cdf) ** 2 + (1.0 - true_cdf) * predicted_cdf**2
    )
    return float(np.trapezoid(expected_squared_distance, grid))

point_crps ¶

point_crps(
    y_true: float,
    samples_pred: ArrayLike,
    x_grid: ArrayLike,
) -> float

Approximate CRPS between a scalar target and empirical predictions.

Parameters:

Name	Type	Description	Default
`y_true`	`float`	Scalar reference target.	required
`samples_pred`	`ArrayLike`	Samples from the predictive scalar distribution.	required
`x_grid`	`ArrayLike`	Finite, strictly increasing integration coordinates.	required

Returns:

Type	Description
`float`	The grid-integrated squared distance between the point and empirical CDFs.

Raises:

Type	Description
`ValueError`	If inputs are empty, non-finite, or malformed.

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

def point_crps(y_true: float, samples_pred: ArrayLike, x_grid: ArrayLike) -> float:
    """Approximate CRPS between a scalar target and empirical predictions.

    Args:
        y_true: Scalar reference target.
        samples_pred: Samples from the predictive scalar distribution.
        x_grid: Finite, strictly increasing integration coordinates.

    Returns:
        The grid-integrated squared distance between the point and empirical CDFs.

    Raises:
        ValueError: If inputs are empty, non-finite, or malformed.
    """
    if not np.isfinite(y_true):
        raise ValueError("y_true must be finite.")
    grid = _evaluation_grid(x_grid)
    predicted_samples = _as_nonempty_vector(samples_pred, name="samples_pred")
    true_cdf = np.where(grid < y_true, 0.0, 1.0)
    predicted_cdf = np.mean(predicted_samples[None, :] <= grid[:, None], axis=1)
    return float(np.trapezoid((true_cdf - predicted_cdf) ** 2, grid))

sample_crps ¶

sample_crps(
    y_true: float, samples_pred: ArrayLike
) -> float

Return the exact CRPS between a scalar target and empirical predictions.

It uses the energy form E|X - y| - E|X - X'| / 2 of the CRPS, with X and X' drawn from the empirical distribution of samples_pred, so no evaluation grid is needed. It equals the integral of the squared distance between the empirical and the target's CDFs over the whole line.

Parameters:

Name	Type	Description	Default
`y_true`	`float`	Scalar reference target.	required
`samples_pred`	`ArrayLike`	Samples from the predictive scalar distribution.	required

Returns:

Type	Description
`float`	The CRPS of the empirical predictive distribution at `y_true`.

Raises:

Type	Description
`ValueError`	If the target is not finite, or the samples are empty, non-finite, or not one-dimensional.

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

def sample_crps(y_true: float, samples_pred: ArrayLike) -> float:
    """Return the exact CRPS between a scalar target and empirical predictions.

    It uses the energy form ``E|X - y| - E|X - X'| / 2`` of the CRPS, with ``X``
    and ``X'`` drawn from the empirical distribution of ``samples_pred``, so
    no evaluation grid is needed. It equals the integral of the squared
    distance between the empirical and the target's CDFs over the whole line.

    Args:
        y_true: Scalar reference target.
        samples_pred: Samples from the predictive scalar distribution.

    Returns:
        The CRPS of the empirical predictive distribution at ``y_true``.

    Raises:
        ValueError: If the target is not finite, or the samples are empty,
            non-finite, or not one-dimensional.
    """
    if not np.isfinite(y_true):
        raise ValueError("y_true must be finite.")
    samples = np.sort(_as_nonempty_vector(samples_pred, name="samples_pred"))
    count = samples.size
    absolute_error = np.mean(np.abs(samples - y_true))
    ranks = 2.0 * np.arange(count) - count + 1.0
    mean_pairwise_distance = 2.0 * np.dot(ranks, samples) / count**2
    return float(absolute_error - mean_pairwise_distance / 2.0)

rmse ¶

rmse(y_true: ArrayLike, y_pred: ArrayLike) -> float

Return the root mean squared error.

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

def rmse(y_true: ArrayLike, y_pred: ArrayLike) -> float:
    """Return the root mean squared error."""
    target, prediction = _matching_vectors(("y_true", y_true), ("y_pred", y_pred))
    return float(np.sqrt(np.mean((target - prediction) ** 2)))

coverage ¶

coverage(
    y_true: ArrayLike, lower: ArrayLike, upper: ArrayLike
) -> float

Return the fraction of targets within inclusive prediction intervals.

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

def coverage(y_true: ArrayLike, lower: ArrayLike, upper: ArrayLike) -> float:
    """Return the fraction of targets within inclusive prediction intervals."""
    target, lower_bound, upper_bound = _matching_vectors(
        ("y_true", y_true), ("lower", lower), ("upper", upper)
    )
    if np.any(lower_bound > upper_bound):
        raise ValueError("lower must not exceed upper.")
    return float(np.mean((lower_bound <= target) & (target <= upper_bound)))

wsu ¶

wsu(
    lower: ArrayLike, upper: ArrayLike, x: ArrayLike
) -> float

Return the normalized weighted spread of uncertainty intervals.

The calculation preserves QMoDeM's established WSU definition while using a generic, strictly increasing coordinate rather than battery-specific time.

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

def wsu(lower: ArrayLike, upper: ArrayLike, x: ArrayLike) -> float:
    """Return the normalized weighted spread of uncertainty intervals.

    The calculation preserves QMoDeM's established WSU definition while using a generic,
    strictly increasing coordinate rather than battery-specific time.
    """
    lower_bound, upper_bound, coordinate = _matching_vectors(
        ("lower", lower), ("upper", upper), ("x", x)
    )
    if coordinate.size < 3:
        raise ValueError("x, lower, and upper must contain at least three values.")
    if np.any(lower_bound > upper_bound):
        raise ValueError("lower must not exceed upper.")
    if np.any(np.diff(coordinate) <= 0):
        raise ValueError("x must be strictly increasing.")

    interval_widths = (upper_bound[2:] + upper_bound[1:-1]) / 2 - (
        lower_bound[2:] + lower_bound[1:-1]
    ) / 2
    period = coordinate[-1] - coordinate[0]
    return float(np.dot(interval_widths, coordinate[1:-1] - coordinate[0]) / period**2)

`stagereg.core.naming`¶

Split ¶

Bases: StrEnum

The partition of data a metric was measured on.

MetricTag ¶

Bases: NamedTuple

The parts of a metric tag.

Attributes:

Name	Type	Description
`stage`	`str`	Name of the stage that produced the metric.
`split`	`Split`	Split the metric was measured on.
`metric`	`str`	Bare name of the metric.

metric_tag ¶

metric_tag(stage: str, split: Split, metric: str) -> str

Build the metric tag for a metric measured in a stage on a split.

Parameters:

Name	Type	Description	Default
`stage`	`str`	Name of the stage that produced the metric.	required
`split`	`Split`	Split the metric was measured on.	required
`metric`	`str`	Bare name of the metric.	required

Returns:

Type	Description
`str`	The tag `stage/split/metric`.

Raises:

Type	Description
`ValueError`	If `stage` or `metric` is empty or contains the separator, or if `split` is not a known split.

Source code in src/stagereg/core/naming.py

def metric_tag(stage: str, split: Split, metric: str) -> str:
    """Build the metric tag for a metric measured in a stage on a split.

    Args:
        stage: Name of the stage that produced the metric.
        split: Split the metric was measured on.
        metric: Bare name of the metric.

    Returns:
        The tag ``stage/split/metric``.

    Raises:
        ValueError: If `stage` or `metric` is empty or contains the
            separator, or if `split` is not a known split.
    """
    _check_segment(stage, "stage")
    _check_segment(metric, "metric")
    return _SEPARATOR.join((stage, Split(split).value, metric))

parse_metric_tag ¶

parse_metric_tag(tag: str) -> MetricTag

Split a metric tag back into its stage, split and metric name.

This is the exact inverse of metric_tag.

Parameters:

Name	Type	Description	Default
`tag`	`str`	A tag of the form `stage/split/metric`.	required

Returns:

Type	Description
`MetricTag`	The stage, split and metric name the tag was built from.

Raises:

Type	Description
`ValueError`	If `tag` does not have exactly three segments, has an empty segment, or names an unknown split.

Source code in src/stagereg/core/naming.py

def parse_metric_tag(tag: str) -> MetricTag:
    """Split a metric tag back into its stage, split and metric name.

    This is the exact inverse of [`metric_tag`][stagereg.core.metric_tag].

    Args:
        tag: A tag of the form ``stage/split/metric``.

    Returns:
        The stage, split and metric name the tag was built from.

    Raises:
        ValueError: If `tag` does not have exactly three segments, has an
            empty segment, or names an unknown split.
    """
    segments = tag.split(_SEPARATOR)
    if len(segments) != 3:
        raise ValueError(
            f"metric tag {tag!r} must have exactly three {_SEPARATOR!r}-separated "
            f"segments, not {len(segments)}."
        )
    stage, split, metric = segments
    _check_segment(stage, "stage")
    _check_segment(metric, "metric")
    try:
        parsed_split = Split(split)
    except ValueError:
        known = ", ".join(repr(member.value) for member in Split)
        raise ValueError(
            f"metric tag {tag!r} names unknown split {split!r}; expected one of "
            f"{known}."
        ) from None
    return MetricTag(stage, parsed_split, metric)

flatten_parameters ¶

flatten_parameters(
    parameters: Mapping[str, Any],
) -> dict[str, Any]

Flatten a nested parameter mapping into parameter paths.

Every non-mapping value, and every empty mapping, is a leaf and is kept unchanged under its key path joined with the separator. Because no key may be empty or contain the separator, two distinct key paths never map onto one parameter path.

Parameters:

Name	Type	Description	Default
`parameters`	`Mapping[str, Any]`	The parameters to flatten, possibly nested.	required

Returns:

Type	Description
`dict[str, Any]`	One entry per leaf, keyed by its parameter path.

Raises:

Type	Description
`ValueError`	If any key, at any depth, is empty or contains the separator.

Source code in src/stagereg/core/naming.py

def flatten_parameters(parameters: Mapping[str, Any]) -> dict[str, Any]:
    """Flatten a nested parameter mapping into parameter paths.

    Every non-mapping value, and every empty mapping, is a leaf and is kept
    unchanged under its key path joined with the separator. Because no key may be empty or contain
    the separator, two distinct key paths never map onto one parameter path.

    Args:
        parameters: The parameters to flatten, possibly nested.

    Returns:
        One entry per leaf, keyed by its parameter path.

    Raises:
        ValueError: If any key, at any depth, is empty or contains the
            separator.
    """

    def walk(mapping: Mapping[str, Any], prefix: str) -> Iterator[tuple[str, Any]]:
        for key, value in mapping.items():
            _check_segment(key, "parameter key")
            path = f"{prefix}{key}"
            if isinstance(value, Mapping) and value:
                yield from walk(value, f"{path}{_SEPARATOR}")
            else:
                yield path, value

    return dict(walk(parameters, ""))

`stagereg.core.protocols`¶

Dataset ¶

Bases: Protocol

An indexable collection of raw examples.

LoaderFactory ¶

Bases: Protocol

Builds a batch loader for a split and epoch.

Optimizer ¶

Bases: Protocol

A functional optimizer with explicit, checkpointable state.

Step ¶

Bases: Protocol

Executes one training or evaluation step.

ValidationStrategy ¶

Bases: Protocol

Evaluates a trained state according to a caller-defined policy.

`stagereg.core.stages`¶

validate_transition ¶

validate_transition(
    current: StageState, next_state: StageState
) -> None

Raise ValueError unless next_state directly follows current.

Source code in src/stagereg/core/stages.py

def validate_transition(current: StageState, next_state: StageState) -> None:
    """Raise ``ValueError`` unless ``next_state`` directly follows ``current``."""
    lifecycle = tuple(StageState)
    current_index = lifecycle.index(current)
    expected_state = (
        lifecycle[current_index + 1] if current_index + 1 < len(lifecycle) else None
    )
    if next_state is not expected_state:
        expected = (
            expected_state.value if expected_state is not None else "no further state"
        )
        raise ValueError(
            f"Cannot transition from {current.value!r} to {next_state.value!r}; "
            f"expected {expected!r}."
        )

TrainingStage ¶

Bases: Protocol

A named unit of training with declared prerequisites and outputs.

`stagereg.core.tracking`¶

Decision `dataclass` ¶

Decision(metric: str, value: float)

The measurement a decision event judged.

Attributes:

Name	Type	Description
`metric`	`str`	Bare name of the monitored metric.
`value`	`float`	The metric's value at the step the decision was made.

TrainingEvent `dataclass` ¶

TrainingEvent(
    name: str,
    stage: str,
    split: Split,
    iteration: int,
    step: int,
    metrics: Mapping[str, float],
    state: TrainingState,
    decision: Decision | None = None,
)

A structured event emitted during staged training.

Attributes:

Name	Type	Description
`name`	`str`	The point in the stage's lifecycle the event marks, such as `epoch_end` or `best_model`.
`stage`	`str`	Name of the stage that emitted the event.
`split`	`Split`	Split the event concerns: the split its metrics were measured on, or, for a decision event, the split of the metric the decision was made on.
`iteration`	`int`	Outer iteration of the staged workflow.
`step`	`int`	Step of the stage the event belongs to.
`metrics`	`Mapping[str, float]`	Metrics measured at this point, keyed by bare metric name. Empty for a decision event.
`state`	`TrainingState`	The live training state.
`decision`	`Decision \| None`	The measurement a decision event judged, or `None` for an event that reports measurements.

EventSink ¶

Bases: Protocol

Consumes training events.

ExperimentTracker ¶

Bases: Protocol

Records parameters, metrics, and artifacts for an experiment.

log_params ¶

log_params(values: Mapping[str, Any]) -> None

Record a run's hyperparameters.

Parameters:

Name	Type	Description	Default
`values`	`Mapping[str, Any]`	A nested mapping whose leaf keys are bare `snake_case` names. A split is a level of nesting of its own (`{"data": {"train": {"samples": ...}}}`), never part of a leaf key (`train_samples`). A tracker that stores flat names flattens the mapping with `flatten_parameters` into parameter paths such as `data/train/samples`.	required

Returns:

Type	Description
`None`	None.

Source code in src/stagereg/core/tracking.py

def log_params(self, values: Mapping[str, Any]) -> None:
    """Record a run's hyperparameters.

    Args:
        values: A nested mapping whose leaf keys are bare
            ``snake_case`` names. A split is a level of nesting of its
            own (``{"data": {"train": {"samples": ...}}}``), never part
            of a leaf key (``train_samples``). A tracker that stores
            flat names flattens the mapping with
            [`flatten_parameters`][stagereg.core.flatten_parameters] into
            parameter paths such as ``data/train/samples``.

    Returns:
        None.
    """
    ...

log_metrics ¶

log_metrics(
    values: Mapping[str, float], *, step: int
) -> None

Record metric values at a step.

Parameters:

Name	Type	Description	Default
`values`	`Mapping[str, float]`	Metric values keyed by the tag to record them under.	required
`step`	`int`	The step the metrics belong to.	required

Returns:

Type	Description
`None`	None.

Source code in src/stagereg/core/tracking.py

def log_metrics(self, values: Mapping[str, float], *, step: int) -> None:
    """Record metric values at a step.

    Args:
        values: Metric values keyed by the tag to record them under.
        step: The step the metrics belong to.

    Returns:
        None.
    """
    ...

log_artifact ¶

log_artifact(name: str, value: Any) -> None

Record an artifact under a name.

Parameters:

Name	Type	Description	Default
`name`	`str`	The name to record the artifact under.	required
`value`	`Any`	The artifact to record.	required

Returns:

Type	Description
`None`	None.

Source code in src/stagereg/core/tracking.py

def log_artifact(self, name: str, value: Any) -> None:
    """Record an artifact under a name.

    Args:
        name: The name to record the artifact under.
        value: The artifact to record.

    Returns:
        None.
    """
    ...

TrackerEventSink `dataclass` ¶

TrackerEventSink(tracker: ExperimentTracker)

An EventSink forwarding event metrics to a tracker.

Each metric is logged under its metric tag, built from the event's stage, the event's split and the metric name, with the event's own step. A decision event carries no metrics, so it adds no series.

The sink calls only ExperimentTracker.log_metrics. Hyperparameters and artifacts stay caller-driven, since neither arrives on a training event.

Attributes:

Name	Type	Description
`tracker`	`ExperimentTracker`	The experiment tracker that receives the metrics.

on_event ¶

on_event(event: TrainingEvent) -> None

Log an event's metrics to the tracker under their metric tags.

Parameters:

Name	Type	Description	Default
`event`	`TrainingEvent`	The training event whose metrics to forward.	required

Returns:

Type	Description
`None`	None.

Source code in src/stagereg/core/tracking.py

def on_event(self, event: TrainingEvent) -> None:
    """Log an event's metrics to the tracker under their metric tags.

    Args:
        event: The training event whose metrics to forward.

    Returns:
        None.
    """
    if not event.metrics:
        return
    self.tracker.log_metrics(
        {
            metric_tag(event.stage, event.split, name): value
            for name, value in event.metrics.items()
        },
        step=event.step,
    )

`stagereg.core.types`¶

Array `module-attribute` ¶

Array = Any

An array of the active backend, such as a jax.Array.

Typed as Any because stagereg.core imports no backend.

PyTree `module-attribute` ¶

PyTree = Any

A nested container of arrays of the active backend, such as model parameters.

Typed as Any because stagereg.core imports no backend.

ParameterRole ¶

Bases: StrEnum

The responsibility of a trainable model component.

StageState ¶

Bases: StrEnum

The ordered lifecycle of a staged probabilistic-regression workflow.

Batch `dataclass` ¶

Batch(
    inputs: PyTree,
    targets: PyTree | None = None,
    sample_weight: Array | None = None,
    metadata: Mapping[str, Any] = dict(),
)

A batch of model inputs, optional targets, weights, and metadata.

Attributes:

Name	Type	Description
`inputs`	`PyTree`	Model inputs, an array or nested container of arrays with a leading batch dimension.
`targets`	`PyTree \| None`	Regression targets aligned with `inputs`. Defaults to `None`, for batches that are only predicted on.
`sample_weight`	`Array \| None`	Weights with a leading batch dimension, one per row, scaling each row's contribution to the loss. Defaults to `None`, which weights every row equally.
`metadata`	`Mapping[str, Any]`	Caller-supplied information about the batch, carried along unchanged. Defaults to an empty mapping.

CheckpointRef `dataclass` ¶

CheckpointRef(
    key: str, metadata: Mapping[str, Any] = dict()
)

An opaque reference to a persisted checkpoint.

Attributes:

Name	Type	Description
`key`	`str`	Key the checkpoint is stored under in a `CheckpointStore`. It identifies the checkpoint and carries no meaning of its own.
`metadata`	`Mapping[str, Any]`	Information about the checkpoint, such as the stage that saved it. Defaults to an empty mapping.

TrainingState `dataclass` ¶

TrainingState(
    model_components: dict[str, Any] = dict(),
    parameter_roles: dict[str, ParameterRole] = dict(),
    frozen_components: frozenset[str] = frozenset(),
    optimizer_states: dict[str, Any] = dict(),
    posterior_state: Any | None = None,
    rng_state: Any | None = None,
    lifecycle_state: StageState = StageState.NEW,
    stage: str | None = None,
    outer_iteration: int = 0,
    data_fingerprint: str | None = None,
    checkpoint_registry: dict[str, CheckpointRef] = dict(),
    metric_history: dict[str, list[float]] = dict(),
)

State shared across one or more explicit training stages.

Attributes:

Name	Type	Description
`model_components`	`dict[str, Any]`	Model objects keyed by component name, as registered with `register_component`. Defaults to an empty mapping.
`parameter_roles`	`dict[str, ParameterRole]`	The `ParameterRole` of each registered component, keyed by component name, so a stage can check the prerequisites an earlier stage left behind. Defaults to an empty mapping.
`frozen_components`	`frozenset[str]`	Names of components that are no longer trained, such as a fitted mean model while a later stage trains the variance. Defaults to an empty set.
`optimizer_states`	`dict[str, Any]`	Optimizer states keyed by name, as registered with `register_optimizer`. Defaults to an empty mapping.
`posterior_state`	`Any \| None`	State of the posterior-approximation stage, set once `lifecycle_state` reaches `StageState.POSTERIOR_READY`. Defaults to `None`.
`rng_state`	`Any \| None`	Random key threaded through training; seed it to make a run reproducible. Defaults to `None`.
`lifecycle_state`	`StageState`	Position of the workflow in the ordered `StageState` lifecycle. Defaults to `StageState.NEW`.
`stage`	`str \| None`	Label of the active stage, the value behind the `active_stage` property. Defaults to `None` outside a runner.
`outer_iteration`	`int`	Zero-based outer iteration: the count of passes through the run's sequence of stages. Defaults to `0`.
`data_fingerprint`	`str \| None`	Identifier of the dataset the state was trained on, meant to detect resuming against different data. Set by the user and not interpreted by the library. Defaults to `None`.
`checkpoint_registry`	`dict[str, CheckpointRef]`	Named references to checkpoints a stage registered, such as its best checkpoint. Defaults to an empty mapping.
`metric_history`	`dict[str, list[float]]`	Recorded values of each metric in order, keyed by metric tag (`stage/split/metric`), the same tags the experiment tracker uses. Defaults to an empty mapping.

active_stage `property` `writable` ¶

active_stage: str | None

Return the explicit active training-stage label.

Returns:

Type	Description
`str \| None`	The active training-stage label, or `None` outside a runner.

register_component ¶

register_component(name: str, component: Any) -> None

Register a model component under name.

Parameters:

Name	Type	Description	Default
`name`	`str`	The key under which `component` is stored in `model_components`.	required
`component`	`Any`	The model component to register.	required

Raises:

Type	Description
`ValueError`	If `name` is already bound to a different object.

Source code in src/stagereg/core/types.py

def register_component(self, name: str, component: Any) -> None:
    """Register a model component under ``name``.

    Args:
        name: The key under which ``component`` is stored in
            ``model_components``.
        component: The model component to register.

    Raises:
        ValueError: If ``name`` is already bound to a different object.
    """
    registered = self.model_components.get(name)
    if name in self.model_components and registered is not component:
        raise ValueError(f"model component {name!r} is already registered.")
    self.model_components[name] = component

register_optimizer ¶

register_optimizer(name: str, optimizer: Any) -> None

Register an optimizer state under name.

Parameters:

Name	Type	Description	Default
`name`	`str`	The key under which `optimizer` is stored in `optimizer_states`.	required
`optimizer`	`Any`	The optimizer state to register.	required

Raises:

Type	Description
`ValueError`	If `name` is already bound to a different object.

Source code in src/stagereg/core/types.py

def register_optimizer(self, name: str, optimizer: Any) -> None:
    """Register an optimizer state under ``name``.

    Args:
        name: The key under which ``optimizer`` is stored in
            ``optimizer_states``.
        optimizer: The optimizer state to register.

    Raises:
        ValueError: If ``name`` is already bound to a different object.
    """
    registered = self.optimizer_states.get(name)
    if name in self.optimizer_states and registered is not optimizer:
        raise ValueError(f"optimizer state {name!r} is already registered.")
    self.optimizer_states[name] = optimizer

record_metric ¶

record_metric(name: str, value: float) -> None

Append value to the metric history recorded under name.

Parameters:

Name	Type	Description	Default
`name`	`str`	The metric name whose history is appended to.	required
`value`	`float`	The observed metric value to append.	required

Source code in src/stagereg/core/types.py

def record_metric(self, name: str, value: float) -> None:
    """Append ``value`` to the metric history recorded under ``name``.

    Args:
        name: The metric name whose history is appended to.
        value: The observed metric value to append.
    """
    self.metric_history.setdefault(name, []).append(value)

StageResult `dataclass` ¶

StageResult(
    state: TrainingState,
    metrics: Mapping[str, float] = dict(),
    loss: float | None = None,
)

The outcome of executing a training stage.

Attributes:

Name	Type	Description
`state`	`TrainingState`	Training state the stage left behind.
`metrics`	`Mapping[str, float]`	Final metric values of the stage, keyed by metric name. Defaults to an empty mapping.
`loss`	`float \| None`	Loss of the stage's last epoch on the `train` split. Defaults to `None` when the stage reports none.

ValidationResult `dataclass` ¶

ValidationResult(
    passed: bool,
    metrics: Mapping[str, float] = dict(),
    message: str | None = None,
)

The outcome of validating a training stage.

Attributes:

Name	Type	Description
`passed`	`bool`	Whether the stage's invariants hold.
`metrics`	`Mapping[str, float]`	Metric values measured during validation, keyed by metric name. Defaults to an empty mapping.
`message`	`str \| None`	Explanation of the outcome, typically of a failure. Defaults to `None`.

stagereg.core¶

stagereg.core.checkpoints¶

Checkpoint dataclass ¶

CheckpointStore ¶

InMemoryCheckpointStore ¶

save ¶

load ¶

exists ¶

stagereg.core.distributions¶

PredictiveDistribution ¶

DistributionHead ¶

Likelihood ¶

DistributionLoss ¶

PredictionLoss ¶

Loss module-attribute ¶

stagereg.core.early_stopping¶

OptimizationMode ¶

EarlyStoppingState dataclass ¶

EarlyStoppingDecision dataclass ¶

should_stop property ¶

EarlyStopper ¶

state property ¶

expects_validation ¶

monitored_metric_name ¶

observe ¶

stagereg.core.losses¶

NegativeLogLikelihoodLoss dataclass ¶

per_example ¶

SquaredErrorLoss dataclass ¶

per_example ¶

add_epsilon ¶

GaussianNLLLoss module-attribute ¶

BetaNLLLoss module-attribute ¶

stagereg.core.metric_registry¶

EvaluationGrid dataclass ¶

PredictionInterval dataclass ¶

EpochPredictionData dataclass ¶

interval ¶

MetricRequirements dataclass ¶

union ¶

EpochMetric ¶

name property ¶

requirements property ¶

__call__ ¶

RootMeanSquaredError dataclass ¶

requirements property ¶

__call__ ¶

NegativeLogLikelihood dataclass ¶

requirements property ¶

__call__ ¶

IntervalCoverage dataclass ¶

requirements property ¶

__call__ ¶

WeightedSpread dataclass ¶

requirements property ¶

__call__ ¶

PointContinuousRankedProbabilityScore dataclass ¶

requirements property ¶

__call__ ¶

SampleContinuousRankedProbabilityScore dataclass ¶

requirements property ¶

__call__ ¶

ContinuousRankedProbabilityScore dataclass ¶

requirements property ¶

__call__ ¶

stagereg.core.metrics¶

cdf ¶

crps ¶

point_crps ¶

sample_crps ¶

rmse ¶

coverage ¶

wsu ¶

stagereg.core.naming¶

Split ¶

MetricTag ¶

metric_tag ¶

parse_metric_tag ¶

flatten_parameters ¶

stagereg.core.protocols¶

`stagereg.core`¶

`stagereg.core.checkpoints`¶

Checkpoint `dataclass` ¶

`stagereg.core.distributions`¶

Loss `module-attribute` ¶

`stagereg.core.early_stopping`¶

EarlyStoppingState `dataclass` ¶

EarlyStoppingDecision `dataclass` ¶

should_stop `property` ¶

state `property` ¶

`stagereg.core.losses`¶

NegativeLogLikelihoodLoss `dataclass` ¶

SquaredErrorLoss `dataclass` ¶

GaussianNLLLoss `module-attribute` ¶

BetaNLLLoss `module-attribute` ¶

`stagereg.core.metric_registry`¶

EvaluationGrid `dataclass` ¶

PredictionInterval `dataclass` ¶

EpochPredictionData `dataclass` ¶

MetricRequirements `dataclass` ¶

name `property` ¶

requirements `property` ¶

call ¶

RootMeanSquaredError `dataclass` ¶

requirements `property` ¶

call ¶

NegativeLogLikelihood `dataclass` ¶

requirements `property` ¶

call ¶

IntervalCoverage `dataclass` ¶

requirements `property` ¶

call ¶

WeightedSpread `dataclass` ¶

requirements `property` ¶

call ¶

PointContinuousRankedProbabilityScore `dataclass` ¶

requirements `property` ¶

call ¶

SampleContinuousRankedProbabilityScore `dataclass` ¶

requirements `property` ¶

call ¶

ContinuousRankedProbabilityScore `dataclass` ¶

requirements `property` ¶

call ¶

`stagereg.core.metrics`¶

`stagereg.core.naming`¶

`stagereg.core.protocols`¶

`stagereg.core.stages`¶

`stagereg.core.tracking`¶

Decision `dataclass` ¶

TrainingEvent `dataclass` ¶

TrackerEventSink `dataclass` ¶

`stagereg.core.types`¶

Array `module-attribute` ¶

PyTree `module-attribute` ¶

Batch `dataclass` ¶

CheckpointRef `dataclass` ¶

TrainingState `dataclass` ¶

active_stage `property` `writable` ¶

StageResult `dataclass` ¶

ValidationResult `dataclass` ¶