substrax 0.1.0__py3-none-any.whl

substrax/__init__.py

@@ -0,0 +1,13 @@
+"""Substrax: shared training and hardware infrastructure for the Avitai JAX stack.
+
+The package sits below calibrax, datarax, artifex and opifex and holds the code those
+packages used to carry separately: device information and placement (``devices``), device
+meshes and sharding strategies (``mesh``), data-parallel training helpers (``spmd``), one
+Orbax checkpoint store (``checkpoint``), training callbacks (``callbacks``) and step-wise
+experiment tracking (``tracking``). Import from the subpackages; this module exposes only
+the version.
+"""
+
+__version__ = "0.1.0"
+
+__all__ = ["__version__"]

substrax/callbacks/__init__.py

@@ -0,0 +1,18 @@
+"""Training callbacks, best-metric tracking and early stopping."""
+
+from substrax.callbacks.base import BaseCallback, CallbackList, TrainerLike, TrainingCallback
+from substrax.callbacks.early_stopping import EarlyStoppingCallback, EarlyStoppingConfig
+from substrax.callbacks.plateau import BestMetricTracker, EarlyStopping, PlateauMode
+
+
+__all__ = [
+    "BaseCallback",
+    "BestMetricTracker",
+    "CallbackList",
+    "EarlyStopping",
+    "EarlyStoppingCallback",
+    "EarlyStoppingConfig",
+    "PlateauMode",
+    "TrainerLike",
+    "TrainingCallback",
+]

substrax/callbacks/base.py

@@ -0,0 +1,158 @@
+"""The training-callback protocol, its no-op base and the dispatching list."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Iterator
+from typing import Any, Protocol, runtime_checkable
+
+from flax import nnx
+
+
+@runtime_checkable
+class TrainerLike(Protocol):
+    """What a callback may ask of the trainer that drives it."""
+
+    @property
+    def model(self) -> nnx.Module:
+        """The model being trained."""
+        ...
+
+
+@runtime_checkable
+class TrainingCallback(Protocol):
+    """Lifecycle hooks a training loop fires; implement any subset via ``BaseCallback``.
+
+    Parameters are positional-only so implementations may name them as they like.
+    """
+
+    def on_train_begin(self, trainer: TrainerLike, /) -> None:
+        """Called once before the first epoch."""
+        ...
+
+    def on_train_end(self, trainer: TrainerLike, /) -> None:
+        """Called once after the last epoch."""
+        ...
+
+    def on_epoch_begin(self, trainer: TrainerLike, epoch: int, /) -> None:
+        """Called before each epoch."""
+        ...
+
+    def on_epoch_end(self, trainer: TrainerLike, epoch: int, logs: dict[str, Any], /) -> None:
+        """Called after each epoch with that epoch's metrics."""
+        ...
+
+    def on_batch_begin(self, trainer: TrainerLike, batch: int, /) -> None:
+        """Called before each batch."""
+        ...
+
+    def on_batch_end(self, trainer: TrainerLike, batch: int, logs: dict[str, Any], /) -> None:
+        """Called after each batch with that batch's metrics."""
+        ...
+
+    def on_validation_begin(self, trainer: TrainerLike, /) -> None:
+        """Called before a validation pass."""
+        ...
+
+    def on_validation_end(self, trainer: TrainerLike, logs: dict[str, Any], /) -> None:
+        """Called after a validation pass with its metrics."""
+        ...
+
+
+class BaseCallback:
+    """No-op implementation of every hook; subclasses override what they need."""
+
+    __slots__ = ()
+
+    def on_train_begin(self, _trainer: TrainerLike) -> None:
+        """Called once before the first epoch."""
+
+    def on_train_end(self, _trainer: TrainerLike) -> None:
+        """Called once after the last epoch."""
+
+    def on_epoch_begin(self, _trainer: TrainerLike, _epoch: int) -> None:
+        """Called before each epoch."""
+
+    def on_epoch_end(self, _trainer: TrainerLike, _epoch: int, _logs: dict[str, Any]) -> None:
+        """Called after each epoch with that epoch's metrics."""
+
+    def on_batch_begin(self, _trainer: TrainerLike, _batch: int) -> None:
+        """Called before each batch."""
+
+    def on_batch_end(self, _trainer: TrainerLike, _batch: int, _logs: dict[str, Any]) -> None:
+        """Called after each batch with that batch's metrics."""
+
+    def on_validation_begin(self, _trainer: TrainerLike) -> None:
+        """Called before a validation pass."""
+
+    def on_validation_end(self, _trainer: TrainerLike, _logs: dict[str, Any]) -> None:
+        """Called after a validation pass with its metrics."""
+
+
+class CallbackList:
+    """Dispatches every hook to its callbacks in registration order."""
+
+    __slots__ = ("_callbacks",)
+
+    def __init__(self, callbacks: Iterable[TrainingCallback] | None = None) -> None:
+        """Wrap an initial set of callbacks.
+
+        Args:
+            callbacks: The callbacks to start with; each receives every hook, in this order.
+        """
+        self._callbacks: list[TrainingCallback] = list(callbacks) if callbacks else []
+
+    def add(self, callback: TrainingCallback) -> None:
+        """Append a callback."""
+        self._callbacks.append(callback)
+
+    def remove(self, callback: TrainingCallback) -> None:
+        """Remove a callback; raises ``ValueError`` if it was never added."""
+        self._callbacks.remove(callback)
+
+    def __len__(self) -> int:
+        """The number of callbacks."""
+        return len(self._callbacks)
+
+    def __iter__(self) -> Iterator[TrainingCallback]:
+        """Iterate over the callbacks in dispatch order."""
+        return iter(self._callbacks)
+
+    def on_train_begin(self, trainer: TrainerLike) -> None:
+        """Dispatch ``on_train_begin``."""
+        for callback in self._callbacks:
+            callback.on_train_begin(trainer)
+
+    def on_train_end(self, trainer: TrainerLike) -> None:
+        """Dispatch ``on_train_end``."""
+        for callback in self._callbacks:
+            callback.on_train_end(trainer)
+
+    def on_epoch_begin(self, trainer: TrainerLike, epoch: int) -> None:
+        """Dispatch ``on_epoch_begin``."""
+        for callback in self._callbacks:
+            callback.on_epoch_begin(trainer, epoch)
+
+    def on_epoch_end(self, trainer: TrainerLike, epoch: int, logs: dict[str, Any]) -> None:
+        """Dispatch ``on_epoch_end``."""
+        for callback in self._callbacks:
+            callback.on_epoch_end(trainer, epoch, logs)
+
+    def on_batch_begin(self, trainer: TrainerLike, batch: int) -> None:
+        """Dispatch ``on_batch_begin``."""
+        for callback in self._callbacks:
+            callback.on_batch_begin(trainer, batch)
+
+    def on_batch_end(self, trainer: TrainerLike, batch: int, logs: dict[str, Any]) -> None:
+        """Dispatch ``on_batch_end``."""
+        for callback in self._callbacks:
+            callback.on_batch_end(trainer, batch, logs)
+
+    def on_validation_begin(self, trainer: TrainerLike) -> None:
+        """Dispatch ``on_validation_begin``."""
+        for callback in self._callbacks:
+            callback.on_validation_begin(trainer)
+
+    def on_validation_end(self, trainer: TrainerLike, logs: dict[str, Any]) -> None:
+        """Dispatch ``on_validation_end``."""
+        for callback in self._callbacks:
+            callback.on_validation_end(trainer, logs)

substrax/callbacks/early_stopping.py

@@ -0,0 +1,116 @@
+"""Early stopping callback driven by epoch logs.
+
+Monitors a metric in the epoch logs and stops training when it stops improving, reaches a
+goal, diverges, or becomes non-finite. The best-so-far and stagnation bookkeeping is
+``BestMetricTracker``; this class adds the log lookup, the thresholds and the epoch record.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any, Literal
+
+from substrax.callbacks.base import BaseCallback, TrainerLike
+from substrax.callbacks.plateau import BestMetricTracker, PlateauMode
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class EarlyStoppingConfig:
+    """Configuration for early stopping.
+
+    Attributes:
+        monitor: Metric name to monitor (e.g., "val_loss", "accuracy").
+        min_delta: Minimum change to qualify as an improvement.
+        patience: Number of epochs with no improvement before stopping.
+        mode: "min" if lower is better, "max" if higher is better.
+        check_finite: If True, stop when metric becomes NaN or Inf.
+        stopping_threshold: Stop immediately when metric reaches this value.
+        divergence_threshold: Stop if metric exceeds this value (min mode only).
+    """
+
+    monitor: str = "val_loss"
+    min_delta: float = 0.0
+    patience: int = 10
+    mode: Literal["min", "max"] = "min"
+    check_finite: bool = True
+    stopping_threshold: float | None = None
+    divergence_threshold: float | None = None
+
+
+class EarlyStoppingCallback(BaseCallback):
+    """Stop training when a monitored metric stops improving."""
+
+    __slots__ = ("_stopped_epoch", "_tracker", "config")
+
+    def __init__(self, config: EarlyStoppingConfig) -> None:
+        """Initialize early stopping callback.
+
+        Args:
+            config: Early stopping configuration.
+        """
+        self.config = config
+        self._tracker = BestMetricTracker(mode=config.mode, min_delta=config.min_delta)
+        self._stopped_epoch: int | None = None
+
+    @property
+    def wait_count(self) -> int:
+        """Epochs since the last improvement."""
+        return self._tracker.num_bad_epochs
+
+    @property
+    def best_score(self) -> float | None:
+        """The best monitored value so far, or ``None`` before the metric first appears."""
+        best = self._tracker.best
+        return None if math.isinf(best) else best
+
+    @property
+    def stopped_epoch(self) -> int | None:
+        """The epoch at which stopping was decided, if it has been."""
+        return self._stopped_epoch
+
+    @property
+    def should_stop(self) -> bool:
+        """Whether training should stop."""
+        return self._stopped_epoch is not None
+
+    def on_epoch_end(self, _trainer: TrainerLike, epoch: int, logs: dict[str, Any]) -> None:
+        """Read the monitored metric, if present, and decide whether to stop.
+
+        Args:
+            _trainer: The trainer instance (unused).
+            epoch: Current epoch number.
+            logs: Dictionary of metrics from this epoch.
+        """
+        value = logs.get(self.config.monitor)
+        if value is None:
+            return
+        current = float(value)
+        if self._stops_immediately(current):
+            self._stopped_epoch = epoch
+            return
+        self._tracker.register(current)
+        if self._tracker.num_bad_epochs >= self.config.patience:
+            self._stopped_epoch = epoch
+
+    def _stops_immediately(self, current: float) -> bool:
+        """Whether ``current`` ends training before patience is considered."""
+        if self.config.check_finite and not math.isfinite(current):
+            return True
+        return self._meets_threshold(current) or self._diverged(current)
+
+    def _meets_threshold(self, current: float) -> bool:
+        """Whether ``current`` reaches the stopping threshold (the training goal)."""
+        threshold = self.config.stopping_threshold
+        if threshold is None:
+            return False
+        if self._tracker.mode is PlateauMode.MIN:
+            return current <= threshold
+        return current >= threshold
+
+    def _diverged(self, current: float) -> bool:
+        """Whether ``current`` exceeds the divergence threshold in min mode."""
+        threshold = self.config.divergence_threshold
+        return (
+            threshold is not None and self._tracker.mode is PlateauMode.MIN and current > threshold
+        )

substrax/callbacks/plateau.py

@@ -0,0 +1,121 @@
+"""Metric-driven early stopping.
+
+Reacts to a stream of per-epoch validation metrics. Framework agnostic (pure-Python
+scalar bookkeeping, no JAX state), so it composes with any training loop.
+
+Semantics follow the established reference: :class:`EarlyStopping` mirrors the
+Keras / Lightning callback (stop after ``patience`` epochs without a ``min_delta``
+improvement). Learning-rate plateau decay is ``optax.contrib.reduce_on_plateau``.
+"""
+
+from __future__ import annotations
+
+from enum import StrEnum
+
+
+class PlateauMode(StrEnum):
+    """Whether a monitored metric improves by decreasing or increasing."""
+
+    MIN = "min"
+    MAX = "max"
+
+
+class BestMetricTracker:
+    """Shared best-so-far + stagnation bookkeeping for the plateau callbacks.
+
+    Tracks the best monitored value and the number of consecutive updates without
+    a ``min_delta`` improvement, in either ``"min"`` or ``"max"`` mode.
+    """
+
+    def __init__(self, *, mode: PlateauMode | str, min_delta: float) -> None:
+        """Initialise the tracker.
+
+        Args:
+            mode: ``"min"`` (lower is better) or ``"max"`` (higher is better).
+            min_delta: Minimum absolute change counted as an improvement.
+
+        Raises:
+            ValueError: If ``min_delta`` is negative.
+        """
+        if min_delta < 0.0:
+            raise ValueError(f"min_delta must be non-negative, got {min_delta}.")
+        self._mode = PlateauMode(mode)
+        self._min_delta = float(min_delta)
+        self._best = float("inf") if self._mode is PlateauMode.MIN else float("-inf")
+        self._num_bad_epochs = 0
+
+    def _is_improvement(self, value: float) -> bool:
+        """Return whether ``value`` beats the best by at least ``min_delta``."""
+        if self._mode is PlateauMode.MIN:
+            return value < self._best - self._min_delta
+        return value > self._best + self._min_delta
+
+    def register(self, value: float) -> bool:
+        """Record ``value``; return ``True`` if it improves on the best so far.
+
+        Updates the best value and resets the stagnation counter on an
+        improvement; otherwise increments the stagnation counter.
+        """
+        if self._is_improvement(value):
+            self._best = float(value)
+            self._num_bad_epochs = 0
+            return True
+        self._num_bad_epochs += 1
+        return False
+
+    @property
+    def mode(self) -> PlateauMode:
+        """The improvement direction."""
+        return self._mode
+
+    @property
+    def best(self) -> float:
+        """The best monitored value seen so far."""
+        return self._best
+
+    @property
+    def num_bad_epochs(self) -> int:
+        """Consecutive updates without a ``min_delta`` improvement."""
+        return self._num_bad_epochs
+
+    def _reset_stagnation(self) -> None:
+        """Clear the stagnation counter (e.g. after an LR reduction)."""
+        self._num_bad_epochs = 0
+
+
+class EarlyStopping(BestMetricTracker):
+    """Signal to stop once a monitored metric stops improving."""
+
+    def __init__(
+        self,
+        *,
+        patience: int,
+        min_delta: float = 0.0,
+        mode: PlateauMode | str = PlateauMode.MIN,
+    ) -> None:
+        """Initialise the stopper.
+
+        Args:
+            patience: Epochs without a ``min_delta`` improvement before stopping.
+            min_delta: Minimum absolute change counted as an improvement.
+            mode: ``"min"`` (lower is better) or ``"max"`` (higher is better).
+
+        Raises:
+            ValueError: If ``patience`` is not positive.
+        """
+        super().__init__(mode=mode, min_delta=min_delta)
+        if patience < 1:
+            raise ValueError(f"patience must be >= 1, got {patience}.")
+        self._patience = patience
+
+    def update(self, value: float) -> bool:
+        """Record the latest metric; return ``True`` if it improved on the best."""
+        return self.register(value)
+
+    @property
+    def should_stop(self) -> bool:
+        """Whether the metric has stagnated for ``patience`` epochs."""
+        return self._num_bad_epochs >= self._patience
+
+
+__all__ = ["BestMetricTracker", "EarlyStopping", "PlateauMode"]

substrax/checkpoint/__init__.py

@@ -0,0 +1,6 @@
+"""The checkpoint store protocol and its Orbax implementation."""
+
+from substrax.checkpoint.checkpoint_store import CheckpointStore, ModelLike, OrbaxCheckpointStore
+
+
+__all__ = ["CheckpointStore", "ModelLike", "OrbaxCheckpointStore"]