pipekit-train

Training pipelines for emulators and amortized inference. A thin orchestration layer over existing training tools (Equinox+Optax ships as the v0.1 reference; Lightning and Keras adapter modules are scaffolded for v0.2 / v0.3). Trained models are first-class pipekit.Operators — they drop into inference pipelines (e.g. pipekit_cycle.NeuralForward for emulator handoff). Master plan reference: Report 11.

The full multi-file design lives in the package: packages/pipekit-train/docs/design/.

Datasets

pipekit_train.dataset

Layer 0 — training datasets.

Carrier-agnostic dataset operators that yield (input, target) pairs and provide a stable content hash.

Classes in this module:

• TrainingDataset — abstract Operator base. Defines the __iter__ + content_hash + with_split contract.
• IterableDataset — generic escape hatch wrapping any iterable with an explicit content hash. Streaming, no random access.
• SimulationDataset — wraps a pipekit_cycle.ForwardModel + a prior operator to generate (params, simulator_output) pairs. Supports both single-step emulation and trajectory rollouts.
• CachedDataset — disk/fsspec-backed cache (zarr in v0.1).

See docs/design/api/datasets.md.

CachedDataset

Bases: TrainingDataset

Disk- or cloud-backed cache around any TrainingDataset.

First epoch hits the source and writes to cache_dir. Subsequent epochs read from the cache. Cache is keyed on source.content_hash() so cache files are reused across runs (and across machines when cache_dir is an fsspec URI) as long as the source's config and seed haven't changed.

Parameters:

Name	Type	Description	Default
source	TrainingDataset	The wrapped TrainingDataset. Must be indexable (provide __len__ + __getitem__) — the cache writes into a fixed-shape zarr store, which requires up-front knowledge of the dataset length. IterableDataset is not cacheable.	required
cache_dir	str | PathLike[str]	Either a local filesystem directory or an fsspec URI (s3://bucket/prefix, gcs://..., memory://..., etc.). fsspec backends require the optional [s3] extra (which pulls in fsspec and s3fs); other backends need their own fsspec drivers (gcsfs, adlfs, …). Created on first write.	required
format	Literal['zarr', 'parquet', 'tfrecord']	Storage format. v0.2 ships "zarr" only; "parquet" and "tfrecord" raise NotImplementedError and are scheduled for v0.3.	'zarr'

Raises:

Type	Description
TypeError	if source doesn't provide __len__ + __getitem__.
NotImplementedError	if format is not "zarr".
ImportError	lazily, if zarr is not installed when the cache is first read or written. fsspec backends also raise ImportError lazily if their driver isn't installed.

Notes

• The cache writes raw NumPy-compatible arrays. Non-array carriers (objects, ragged structures) must be converted upfront — CachedDataset is for array-shaped pairs.
• The cache is content-addressed, not split-addressed. Calling cached.with_split("val") returns a CachedDataset wrapping source.with_split("val"); the cache for that split lives at a different path because its content hash differs.
• To force re-cache, delete the cache directory or change a source-config field.

Source code in packages/pipekit-train/src/pipekit_train/dataset.py

class CachedDataset(TrainingDataset):
    """Disk- or cloud-backed cache around any `TrainingDataset`.

    First epoch hits the source and writes to ``cache_dir``. Subsequent
    epochs read from the cache. Cache is keyed on ``source.content_hash()``
    so cache files are reused across runs (and across machines when
    ``cache_dir`` is an fsspec URI) as long as the source's config and
    seed haven't changed.

    Args:
        source: The wrapped `TrainingDataset`. Must be indexable
            (provide ``__len__`` + ``__getitem__``) — the cache writes
            into a fixed-shape zarr store, which requires up-front
            knowledge of the dataset length. `IterableDataset` is not
            cacheable.
        cache_dir: Either a local filesystem directory **or** an
            fsspec URI (``s3://bucket/prefix``, ``gcs://...``,
            ``memory://...``, etc.). fsspec backends require the
            optional ``[s3]`` extra (which pulls in ``fsspec`` and
            ``s3fs``); other backends need their own fsspec drivers
            (``gcsfs``, ``adlfs``, …). Created on first write.
        format: Storage format. v0.2 ships ``"zarr"`` only;
            ``"parquet"`` and ``"tfrecord"`` raise NotImplementedError
            and are scheduled for v0.3.

    Raises:
        TypeError: if ``source`` doesn't provide ``__len__`` + ``__getitem__``.
        NotImplementedError: if ``format`` is not ``"zarr"``.
        ImportError: lazily, if ``zarr`` is not installed when the
            cache is first read or written. fsspec backends also raise
            ImportError lazily if their driver isn't installed.

    Notes:
        - The cache writes raw NumPy-compatible arrays. Non-array
          carriers (objects, ragged structures) must be converted
          upfront — `CachedDataset` is for array-shaped pairs.
        - The cache is content-addressed, not split-addressed. Calling
          ``cached.with_split("val")`` returns a `CachedDataset`
          wrapping ``source.with_split("val")``; the cache for that
          split lives at a different path because its content hash
          differs.
        - To force re-cache, delete the cache directory or change a
          source-config field.
    """

    forbid_in_yaml: ClassVar[bool] = True
    __config_mixin_auto__ = False

    def __init__(
        self,
        source: TrainingDataset,
        cache_dir: str | os.PathLike[str],
        format: Literal["zarr", "parquet", "tfrecord"] = "zarr",
    ) -> None:
        if not isinstance(source, TrainingDataset):
            raise TypeError(
                "CachedDataset.source must be a TrainingDataset; got "
                f"{type(source).__name__}."
            )
        if not (hasattr(source, "__len__") and hasattr(source, "__getitem__")):
            raise TypeError(
                "CachedDataset.source must be indexable (provide __len__ + "
                f"__getitem__); got {type(source).__name__}. Wrap streaming "
                "sources differently or materialise them first."
            )
        if format != "zarr":
            raise NotImplementedError(
                f"CachedDataset format={format!r} is scheduled for v0.3. "
                "Use format='zarr'; see docs/design/api/datasets.md."
            )
        super().__init__(seed=source.seed, split=cast(Split, source.split))
        self.source = source
        self.cache_dir = str(cache_dir)
        self.format: Literal["zarr", "parquet", "tfrecord"] = format

    def with_split(self, split: Split) -> CachedDataset:
        return CachedDataset(
            source=self.source.with_split(split),
            cache_dir=self.cache_dir,
            format=self.format,
        )

    def content_hash(self) -> str:
        # Inherits source.content_hash() — the cache is a materialisation
        # of the same logical contents.
        return self.source.content_hash()

    # --- cache layout --------------------------------------------------

    @staticmethod
    def _is_fsspec_uri(cache_dir: str) -> bool:
        """Detect an fsspec URI (anything with a ``<scheme>://`` prefix).

        Local paths return False even when prefixed with ``file://``
        (we still route those through the local-store path; zarr's
        LocalStore handles both shapes).
        """
        # ``re`` is in the stdlib; safer than parsing manually.
        import re

        match = re.match(r"^([a-zA-Z][a-zA-Z0-9+\-.]*)://", cache_dir)
        if match is None:
            return False
        scheme = match.group(1)
        return scheme != "file"

    def _cache_path(self) -> str:
        """Cache path for the current content hash.

        Returns a string suitable for either ``os.path``-style local
        operations OR fsspec URI routing — the latter case never
        touches `os.path.join` (which would mangle the scheme).
        """
        h = self.content_hash()
        if self._is_fsspec_uri(self.cache_dir):
            # fsspec URIs use '/' separators universally regardless of OS.
            base = self.cache_dir.rstrip("/")
            return f"{base}/{h}"
        return os.path.join(self.cache_dir, h)

    def _open_zarr_root(self, mode: str) -> Any:
        """Open the zarr root for the cache, picking the right store.

        Local paths use zarr's default open (which auto-picks
        `LocalStore`). fsspec URIs route through `zarr.storage.FsspecStore.from_url`,
        which requires the relevant fsspec driver (`fsspec` for memory,
        `s3fs` for s3, etc.) to be importable. ImportError is raised
        lazily by zarr/fsspec if the driver isn't installed.
        """
        import zarr  # ty: ignore[unresolved-import]

        path = self._cache_path()
        if self._is_fsspec_uri(self.cache_dir):
            store = zarr.storage.FsspecStore.from_url(path, read_only=(mode == "r"))
            return cast(Any, zarr.open(store, mode=mode))
        return cast(Any, zarr.open(path, mode=mode))

    def _materialised(self) -> bool:
        """Check whether the cache has already been written."""
        try:
            import zarr  # noqa: F401  # ty: ignore[unresolved-import]
        except ImportError:
            return False
        try:
            root = self._open_zarr_root(mode="r")
        except (FileNotFoundError, KeyError, ValueError):
            return False
        return "x" in root and "y" in root

    def _materialise(self) -> None:
        """Write the source into the zarr cache (first epoch).

        Works for both local paths and fsspec URIs — the zarr store
        backend is chosen in `_open_zarr_root`.
        """
        import numpy as np

        # `self.source` is typed as TrainingDataset; the indexability
        # guard in __init__ means it also implements __len__/__getitem__.
        # ty can't combine the runtime guard with the static type, so
        # we cast at the boundary.
        source = cast(Any, self.source)
        n = len(source)
        if n == 0:
            raise ValueError(
                "CachedDataset cannot materialise an empty source — "
                "we need at least one sample to infer the array "
                "shapes / dtypes for the zarr store. Check the source "
                "dataset's split sizes or content_hash."
            )
        first_x, first_y = source[0]
        first_x = np.asarray(first_x)
        first_y = np.asarray(first_y)

        # Local paths need the directory created on disk; fsspec
        # backends create implicit "directories" on key write.
        if not self._is_fsspec_uri(self.cache_dir):
            os.makedirs(self._cache_path(), exist_ok=True)

        root = self._open_zarr_root(mode="w")
        x_arr = root.create_array("x", shape=(n, *first_x.shape), dtype=first_x.dtype)
        y_arr = root.create_array("y", shape=(n, *first_y.shape), dtype=first_y.dtype)
        x_arr[0] = first_x
        y_arr[0] = first_y
        for i in range(1, n):
            xi, yi = source[i]
            x_arr[i] = np.asarray(xi)
            y_arr[i] = np.asarray(yi)

    def _open_cache(self) -> tuple[Any, Any]:
        root = self._open_zarr_root(mode="r")
        return root["x"], root["y"]

    # --- iteration / indexing -----------------------------------------

    def __len__(self) -> int:
        return len(cast(Any, self.source))

    def __getitem__(self, idx: int) -> tuple[Any, Any]:
        if not self._materialised():
            self._materialise()
        x_arr, y_arr = self._open_cache()
        return x_arr[idx], y_arr[idx]

    def __iter__(self) -> Iterator[tuple[Any, Any]]:
        if not self._materialised():
            self._materialise()
        x_arr, y_arr = self._open_cache()
        n = x_arr.shape[0]
        for i in range(n):
            yield x_arr[i], y_arr[i]

    # --- explicit invalidation ----------------------------------------

    def invalidate(self) -> None:
        """Delete the cache for the current content hash.

        After calling this, the next iteration re-materialises from
        the source. Useful when the cache becomes corrupt or when
        you want to force a re-bake without changing source config.
        Works for both local paths and fsspec URIs.
        """
        path = self._cache_path()
        if self._is_fsspec_uri(self.cache_dir):
            import fsspec

            fs, fs_path = fsspec.core.url_to_fs(path)
            if fs.exists(fs_path):
                fs.rm(fs_path, recursive=True)
        else:
            import shutil

            if os.path.exists(path):
                shutil.rmtree(path)

    def get_config(self) -> dict[str, Any]:
        return {
            "source": type(self.source).__qualname__,
            "source_hash": self.source.content_hash(),
            "cache_dir": self.cache_dir,
            "format": self.format,
        }

invalidate()

Delete the cache for the current content hash.

After calling this, the next iteration re-materialises from the source. Useful when the cache becomes corrupt or when you want to force a re-bake without changing source config. Works for both local paths and fsspec URIs.

Source code in packages/pipekit-train/src/pipekit_train/dataset.py

def invalidate(self) -> None:
    """Delete the cache for the current content hash.

    After calling this, the next iteration re-materialises from
    the source. Useful when the cache becomes corrupt or when
    you want to force a re-bake without changing source config.
    Works for both local paths and fsspec URIs.
    """
    path = self._cache_path()
    if self._is_fsspec_uri(self.cache_dir):
        import fsspec

        fs, fs_path = fsspec.core.url_to_fs(path)
        if fs.exists(fs_path):
            fs.rm(fs_path, recursive=True)
    else:
        import shutil

        if os.path.exists(path):
            shutil.rmtree(path)

IterableDataset

Bases: TrainingDataset

Generic escape hatch — wrap any iterable as a TrainingDataset.

Useful when the data source doesn't fit the catalog / simulation patterns: an existing PyTorch dataset, a generator over remote files, a user-built sequence of synthetic batches.

Parameters:

Name	Type	Description	Default
source	Iterable[tuple[Any, Any]] | Callable[[], Iterable[tuple[Any, Any]]]	Either an iterable of (input, target) pairs, or a zero-arg callable returning such an iterable (callable form is required when the iterable is single-pass — the trainer needs to iterate multiple epochs).	required
content_hash	str	User-supplied content hash. Required, because we can't introspect source deterministically.	required
seed	int	PRNG seed (unused by the base wrapper; surfaced for consistency with sampling subclasses).	0
split	Split	"train" (default), "val", or "test". The base wrapper does NOT partition; it's the user's responsibility to pass a split-specific source.	'train'

Because the source is an opaque user closure, this class sets forbid_in_yaml = True — the YAML round-trip is debug-only.

Equinox-adapter note: IterableDatasets aren't indexable, so the adapter falls back from Grain's MapDataset (indexed) to plain Python iteration with manual batching.

Source code in packages/pipekit-train/src/pipekit_train/dataset.py

class IterableDataset(TrainingDataset):
    """Generic escape hatch — wrap any iterable as a `TrainingDataset`.

    Useful when the data source doesn't fit the catalog / simulation
    patterns: an existing PyTorch dataset, a generator over remote
    files, a user-built sequence of synthetic batches.

    Args:
        source: Either an iterable of ``(input, target)`` pairs, or a
            zero-arg callable returning such an iterable (callable form
            is required when the iterable is single-pass — the trainer
            needs to iterate multiple epochs).
        content_hash: User-supplied content hash. Required, because we
            can't introspect ``source`` deterministically.
        seed: PRNG seed (unused by the base wrapper; surfaced for
            consistency with sampling subclasses).
        split: ``"train"`` (default), ``"val"``, or ``"test"``. The
            base wrapper does NOT partition; it's the user's
            responsibility to pass a split-specific source.

    Because the source is an opaque user closure, this class sets
    ``forbid_in_yaml = True`` — the YAML round-trip is debug-only.

    Equinox-adapter note: IterableDatasets aren't indexable, so the
    adapter falls back from Grain's `MapDataset` (indexed) to plain
    Python iteration with manual batching.
    """

    forbid_in_yaml: ClassVar[bool] = True
    __config_mixin_auto__ = False

    def __init__(
        self,
        source: Iterable[tuple[Any, Any]] | Callable[[], Iterable[tuple[Any, Any]]],
        content_hash: str,
        seed: int = 0,
        split: Split = "train",
    ) -> None:
        super().__init__(seed=seed, split=split)
        if not (callable(source) or hasattr(source, "__iter__")):
            raise TypeError(
                "IterableDataset.source must be an iterable or a zero-arg "
                f"callable returning one; got {type(source).__name__}."
            )
        if not isinstance(content_hash, str) or not content_hash:
            raise ValueError(
                "IterableDataset.content_hash must be a non-empty string — "
                "the framework can't derive a hash for opaque iterables."
            )
        self.source = source
        self._content_hash = content_hash

    def __iter__(self) -> Iterator[tuple[Any, Any]]:
        # ty can't narrow the union; cast both branches via Any.
        raw_source = cast(Any, self.source)
        source = raw_source() if callable(raw_source) else raw_source
        yield from source

    def content_hash(self) -> str:
        return hash_config(
            "IterableDataset",
            {"hash": self._content_hash, "seed": self.seed, "split": self.split},
        )

    def get_config(self) -> dict[str, Any]:
        return {
            "content_hash": self._content_hash,
            "seed": self.seed,
            "split": self.split,
        }

SimulationDataset

Bases: TrainingDataset

Yields (parameters, simulator_output) pairs.

Used for emulator training and amortized inference. The dataset samples a prior, evaluates the forward model (single step or trajectory), and optionally pushes the output through an observation operator (for SBI where the network learns to invert observations, not raw model states).

Parameters:

Name	Type	Description	Default
forward_model	Any	Anything satisfying pipekit_cycle.ForwardModel — exposes step(state, dt) and dt. Workspace-internal package; resolves through the [cycle] extra at runtime.	required
prior	Operator	Zero-arg Operator returning a parameter sample.	required
n_samples	int	Total samples per epoch.	required
cycle	Operator | None	Optional pipekit_cycle.Cycle for trajectory rollouts. When None, the dataset emits single-step pairs.	None
obs_op	Any | None	Optional pipekit_cycle.ObservationOperator. When set, the output is passed through it before yielding (used for SBI training where the target is observed-space outputs rather than latent state).	None
seed	int	PRNG seed for the prior. Each sample uses seed + split_offset + i so splits don't overlap.	0
split	Split	"train", "val", or "test". Partitions the seed space (80/10/10 by default).	'train'

The content hash includes the forward-model state_signature (when available) and dt, the prior config, n_samples, seed, and split — anything that affects the data.

Source code in packages/pipekit-train/src/pipekit_train/dataset.py

class SimulationDataset(TrainingDataset):
    """Yields ``(parameters, simulator_output)`` pairs.

    Used for emulator training and amortized inference. The dataset
    samples a prior, evaluates the forward model (single step or
    trajectory), and optionally pushes the output through an
    observation operator (for SBI where the network learns to invert
    *observations*, not raw model states).

    Args:
        forward_model: Anything satisfying
            ``pipekit_cycle.ForwardModel`` — exposes ``step(state, dt)``
            and ``dt``. Workspace-internal package; resolves through the
            ``[cycle]`` extra at runtime.
        prior: Zero-arg `Operator` returning a parameter sample.
        n_samples: Total samples per epoch.
        cycle: Optional ``pipekit_cycle.Cycle`` for trajectory rollouts.
            When ``None``, the dataset emits single-step pairs.
        obs_op: Optional ``pipekit_cycle.ObservationOperator``. When
            set, the output is passed through it before yielding
            (used for SBI training where the target is observed-space
            outputs rather than latent state).
        seed: PRNG seed for the prior. Each sample uses
            ``seed + split_offset + i`` so splits don't overlap.
        split: ``"train"``, ``"val"``, or ``"test"``. Partitions the
            seed space (80/10/10 by default).

    The content hash includes the forward-model ``state_signature``
    (when available) and ``dt``, the prior config, ``n_samples``,
    ``seed``, and ``split`` — anything that affects the data.
    """

    forbid_in_yaml: ClassVar[bool] = True  # forward_model often carries closures
    __config_mixin_auto__ = False

    # Split → (numerator_start, numerator_end) over a 10ths denominator
    # so the integer split sizes are exact (80 / 10 / 10).
    _SPLIT_TENTHS: ClassVar[dict[Split, tuple[int, int]]] = {
        "train": (0, 8),
        "val": (8, 9),
        "test": (9, 10),
    }

    def __init__(
        self,
        forward_model: Any,
        prior: Operator,
        n_samples: int,
        cycle: Operator | None = None,
        obs_op: Any | None = None,
        seed: int = 0,
        split: Split = "train",
    ) -> None:
        super().__init__(seed=seed, split=split)
        for name in ("step", "dt"):
            if not hasattr(forward_model, name):
                raise TypeError(
                    f"SimulationDataset.forward_model must satisfy "
                    f"pipekit_cycle.ForwardModel (missing `{name}`). Got "
                    f"{type(forward_model).__name__}."
                )
        if not isinstance(prior, Operator):
            raise TypeError(
                "SimulationDataset.prior must be a pipekit.Operator; got "
                f"{type(prior).__name__}."
            )
        if n_samples <= 0:
            raise ValueError(
                f"SimulationDataset.n_samples must be positive; got {n_samples}."
            )
        self.forward_model = forward_model
        self.prior = prior
        self.n_samples = int(n_samples)
        self.cycle = cycle
        self.obs_op = obs_op

    # --- iteration / indexing -------------------------------------------

    def __len__(self) -> int:
        start_t, end_t = self._SPLIT_TENTHS[cast(Split, self.split)]
        end = self.n_samples * end_t // 10
        start = self.n_samples * start_t // 10
        return end - start

    def __getitem__(self, idx: int) -> tuple[Any, Any]:
        n = len(self)
        if idx < 0:
            idx += n
        if not 0 <= idx < n:
            raise IndexError(idx)
        # Per-sample deterministic seed so splits don't overlap.
        sample_seed = self.seed + self._split_offset() + idx
        # The prior's call protocol: takes an optional seed kwarg if it
        # accepts one; otherwise just call.
        try:
            params = self.prior(seed=sample_seed)
        except TypeError:
            params = self.prior()
        if self.cycle is not None:
            output = self.cycle(params)
        else:
            output = self.forward_model.step(params, self.forward_model.dt)
        if self.obs_op is not None:
            output = self.obs_op(output)
        return params, output

    def __iter__(self) -> Iterator[tuple[Any, Any]]:
        for i in range(len(self)):
            yield self[i]

    def _split_offset(self) -> int:
        """Stable per-split offset into the seed space."""
        start_t, _ = self._SPLIT_TENTHS[cast(Split, self.split)]
        return self.n_samples * start_t // 10

    # --- hashing / config ----------------------------------------------

    def content_hash(self) -> str:
        fm_sig = getattr(self.forward_model, "state_signature", None)
        fm_dt = getattr(self.forward_model, "dt", None)
        return hash_config(
            "SimulationDataset",
            {
                "forward_model": {
                    "class": type(self.forward_model).__qualname__,
                    "state_signature": repr(fm_sig),
                    "dt": fm_dt,
                },
                "prior": {
                    "class": type(self.prior).__qualname__,
                    "config": self.prior.get_config(),
                },
                "n_samples": self.n_samples,
                "cycle": (
                    type(self.cycle).__qualname__ if self.cycle is not None else None
                ),
                "obs_op": (
                    type(self.obs_op).__qualname__ if self.obs_op is not None else None
                ),
                "seed": self.seed,
                "split": self.split,
            },
        )

    def get_config(self) -> dict[str, Any]:
        return {
            "forward_model": type(self.forward_model).__qualname__,
            "prior": type(self.prior).__qualname__,
            "n_samples": self.n_samples,
            "cycle": (
                type(self.cycle).__qualname__ if self.cycle is not None else None
            ),
            "obs_op": (
                type(self.obs_op).__qualname__ if self.obs_op is not None else None
            ),
            "seed": self.seed,
            "split": self.split,
        }

TrainingDataset

Bases: Operator

Abstract base — yields (input, target) pairs.

Subclasses implement __iter__ and content_hash. The base provides:

• pipekit.Operator semantics — get_config / dumps / loads.
• with_split for deterministic train/val/test split derivation.
• A default content_hash over (type, get_config()) that subclasses can extend or override.

The training surface is the iterator, not _apply. Calling a TrainingDataset like dataset() is undefined and raises TypeError — datasets compose into TrainingLoop via the dataset= constructor argument, not via Sequential.

Attributes:

Name	Type	Description
seed		PRNG seed used by sampling subclasses. Default 0.
split		"train" (default), "val", or "test".

Source code in packages/pipekit-train/src/pipekit_train/dataset.py

class TrainingDataset(Operator):
    """Abstract base — yields ``(input, target)`` pairs.

    Subclasses implement ``__iter__`` and ``content_hash``. The base
    provides:

    - `pipekit.Operator` semantics — get_config / dumps / loads.
    - `with_split` for deterministic train/val/test split derivation.
    - A default ``content_hash`` over ``(type, get_config())`` that
      subclasses can extend or override.

    The training surface is the **iterator**, not ``_apply``. Calling
    a `TrainingDataset` like ``dataset()`` is undefined and raises
    ``TypeError`` — datasets compose into `TrainingLoop` via the
    ``dataset=`` constructor argument, not via `Sequential`.

    Attributes:
        seed: PRNG seed used by sampling subclasses. Default ``0``.
        split: ``"train"`` (default), ``"val"``, or ``"test"``.
    """

    # Datasets serialise their config but the data source typically
    # references external state (catalog URI, forward model with closures).
    # Mark the class so YAML loaders can decide policy per subclass.
    forbid_in_yaml: ClassVar[bool] = False

    def __init__(self, seed: int = 0, split: Split = "train") -> None:
        if split not in ("train", "val", "test"):
            raise ValueError(
                f"split must be one of 'train', 'val', 'test'; got {split!r}."
            )
        self.seed = int(seed)
        self.split = split

    def __iter__(self) -> Iterator[tuple[Any, Any]]:
        raise NotImplementedError(
            f"{type(self).__name__} must implement __iter__ to yield "
            "(input, target) pairs."
        )

    def _apply(self, *args: Any, **kwargs: Any) -> Any:
        raise TypeError(
            f"{type(self).__name__} is iterated, not called. Use "
            "`for x, y in dataset:` or pass it to `TrainingLoop(dataset=...)`."
        )

    def content_hash(self) -> str:
        """Stable identifier for this dataset's contents.

        Default folds ``(type, get_config())`` into a sha256. Subclasses
        with external state (catalog URIs, forward-model signatures)
        override to fold those in as well.
        """
        return hash_config(type(self).__qualname__, self.get_config())

    def with_split(self, split: Split) -> TrainingDataset:
        """Return a shallow clone of this dataset with a different split.

        Same ``seed``; only the split label changes. Subclasses that
        carry mutable state should override.
        """
        clone = self.__class__.__new__(self.__class__)
        clone.__dict__.update(self.__dict__)
        clone.split = split
        return clone

content_hash()

Stable identifier for this dataset's contents.

Default folds (type, get_config()) into a sha256. Subclasses with external state (catalog URIs, forward-model signatures) override to fold those in as well.

Source code in packages/pipekit-train/src/pipekit_train/dataset.py

def content_hash(self) -> str:
    """Stable identifier for this dataset's contents.

    Default folds ``(type, get_config())`` into a sha256. Subclasses
    with external state (catalog URIs, forward-model signatures)
    override to fold those in as well.
    """
    return hash_config(type(self).__qualname__, self.get_config())

with_split(split)

Return a shallow clone of this dataset with a different split.

Same seed; only the split label changes. Subclasses that carry mutable state should override.

Source code in packages/pipekit-train/src/pipekit_train/dataset.py

def with_split(self, split: Split) -> TrainingDataset:
    """Return a shallow clone of this dataset with a different split.

    Same ``seed``; only the split label changes. Subclasses that
    carry mutable state should override.
    """
    clone = self.__class__.__new__(self.__class__)
    clone.__dict__.update(self.__dict__)
    clone.split = split
    return clone

hash_config(*parts)

sha256 over parts, each stable-JSON-encoded and NUL-terminated.

The shared content-hash framing for TrainingDataset.content_hash implementations (used by SimulationDataset, CachedDataset, and geotoolz-patcher's geopatcher.integrations.train.XarrayWindowDataset).

Source code in packages/pipekit-train/src/pipekit_train/dataset.py

def hash_config(*parts: Any) -> str:
    """sha256 over ``parts``, each stable-JSON-encoded and NUL-terminated.

    The shared content-hash framing for `TrainingDataset.content_hash`
    implementations (used by `SimulationDataset`, `CachedDataset`, and
    geotoolz-patcher's `geopatcher.integrations.train.XarrayWindowDataset`).
    """
    chunks: list[bytes | str] = []
    for p in parts:
        chunks += [stable_json(p), b"\x00"]
    return sha256_hex(*chunks)

Losses

pipekit_train.loss

Loss protocol + common loss operators.

The carrier-agnostic loss surface from the v0.1 design.

• Loss (Protocol) — runtime-checkable; signature (predicted, target) -> float | tuple[float, dict].
• MSE — mean squared error. Carrier-agnostic via duck typing on -, **, .mean; works for both numpy and JAX arrays.
• NLL — assumes predicted has .log_prob(target). Distribution- shaped outputs (Equinox distributions, numpyro distributions, custom density heads). For Gaussian-NLL on raw (mean, log_var) tuples, wrap them in a small density head that exposes .log_prob.
• KL — Kullback-Leibler divergence. Prefers predicted.kl_divergence (target) if available; otherwise falls back to -(target.log_prob(samples) - predicted.log_prob(samples)).mean() with samples drawn from predicted.
• Composite — weighted sum of Loss components, emitting the per-component values as auxiliary metrics.

See docs/design/api/primitives.md.

Composite

Bases: Operator

Weighted sum of Loss components.

The aggregate loss is sum(w_i * L_i(predicted, target)). Each component's scalar value is surfaced as an auxiliary metric so the adapter logs them individually.

Parameters:

Name	Type	Description	Default
components	tuple[tuple[float, Loss], ...]	Tuple of (weight, loss) pairs. weight is a Python float; loss is any Loss-shaped callable.	required

Raises:

Type	Description
ValueError	if components is empty.

Returns from __call__: (total, {"<name>_<i>": value, ..., "total": total}).

Source code in packages/pipekit-train/src/pipekit_train/loss.py

class Composite(Operator):
    """Weighted sum of `Loss` components.

    The aggregate loss is ``sum(w_i * L_i(predicted, target))``. Each
    component's scalar value is surfaced as an auxiliary metric so
    the adapter logs them individually.

    Args:
        components: Tuple of ``(weight, loss)`` pairs. ``weight`` is
            a Python float; ``loss`` is any `Loss`-shaped callable.

    Raises:
        ValueError: if ``components`` is empty.

    Returns from ``__call__``:
        ``(total, {"<name>_<i>": value, ..., "total": total})``.
    """

    __config_mixin_auto__ = False

    def __init__(
        self,
        components: tuple[tuple[float, Loss], ...],
    ) -> None:
        if not components:
            raise ValueError("Composite requires at least one component.")
        self.components = tuple(components)

    def _apply(
        self,
        predicted: Any,
        target: Any,
    ) -> tuple[Any, dict[str, Any]]:
        aux: dict[str, Any] = {}
        total = None
        for i, (weight, loss) in enumerate(self.components):
            result = loss(predicted, target)
            value = result[0] if isinstance(result, tuple) else result
            name = f"{type(loss).__name__}_{i}"
            aux[name] = value
            weighted = weight * value
            total = weighted if total is None else total + weighted
        aux["total"] = total
        return total, aux

    def get_config(self) -> dict[str, Any]:
        # The components tuple holds Loss callables that may be
        # arbitrary operators — surface their class names and weights.
        return {
            "components": [
                {"weight": float(w), "loss_class": type(loss).__name__}
                for w, loss in self.components
            ]
        }

KL

Bases: Operator

Kullback-Leibler divergence between two distributions.

Contract

predicted and target must both expose either .kl_divergence(other) (fast analytic path) or .log_prob(samples) + .sample(key) (sample-based fallback). The fast path is tried first.

Returns from __call__: (loss, {"kl": loss}).

Note

The sample-based fallback requires a PRNG key. v0.1 ships the analytic path only; sample-based KL lands in v0.1.1 when the Equinox adapter exposes the per-batch RNG.

Source code in packages/pipekit-train/src/pipekit_train/loss.py

class KL(Operator):
    """Kullback-Leibler divergence between two distributions.

    Contract:
        ``predicted`` and ``target`` must both expose either
        ``.kl_divergence(other)`` (fast analytic path) or
        ``.log_prob(samples)`` + ``.sample(key)`` (sample-based
        fallback). The fast path is tried first.

    Returns from ``__call__``:
        ``(loss, {"kl": loss})``.

    Note:
        The sample-based fallback requires a PRNG key. v0.1 ships the
        analytic path only; sample-based KL lands in v0.1.1 when the
        Equinox adapter exposes the per-batch RNG.
    """

    def _apply(
        self,
        predicted: Any,
        target: Any,
    ) -> tuple[Any, dict[str, Any]]:
        if hasattr(predicted, "kl_divergence"):
            loss = predicted.kl_divergence(target).mean()
            return loss, {"kl": loss}
        raise TypeError(
            "KL requires `predicted` to expose `.kl_divergence(target)`. "
            "Sample-based KL fallback is scheduled for v0.1.1; see "
            "docs/design/api/primitives.md."
        )

Loss

Bases: Protocol

Loss function — per-batch scalar called by the adapter.

Implementations may return either a scalar loss directly, or a (scalar, aux_metrics_dict) tuple. The adapter distinguishes by introspection.

See docs/design/api/primitives.md and ADR D4 in docs/design/decisions.md.

Source code in packages/pipekit-train/src/pipekit_train/loss.py

@runtime_checkable
class Loss(Protocol):
    """Loss function — per-batch scalar called by the adapter.

    Implementations may return either a scalar loss directly, or a
    ``(scalar, aux_metrics_dict)`` tuple. The adapter distinguishes
    by introspection.

    See ``docs/design/api/primitives.md`` and ADR D4 in
    ``docs/design/decisions.md``.
    """

    def __call__(
        self,
        predicted: Any,
        target: Any,
    ) -> float | tuple[float, dict[str, float]]: ...

MSE

Bases: Operator

Mean squared error loss.

Parameters:

Name	Type	Description	Default
reduction	str	"mean" (default), "sum", or "none". "none" returns the elementwise squared error without reduction — the adapter is then responsible for reducing.	'mean'

Returns from __call__: (loss, {"mse": loss}) so adapters see a consistent (scalar, aux) shape.

Source code in packages/pipekit-train/src/pipekit_train/loss.py

class MSE(Operator):
    """Mean squared error loss.

    Args:
        reduction: ``"mean"`` (default), ``"sum"``, or ``"none"``.
            ``"none"`` returns the elementwise squared error without
            reduction — the adapter is then responsible for reducing.

    Returns from ``__call__``:
        ``(loss, {"mse": loss})`` so adapters see a consistent
        ``(scalar, aux)`` shape.
    """

    def __init__(self, reduction: str = "mean") -> None:
        if reduction not in {"mean", "sum", "none"}:
            raise ValueError(
                f"MSE.reduction must be one of 'mean', 'sum', 'none'; "
                f"got {reduction!r}."
            )
        self.reduction = reduction

    def _apply(
        self,
        predicted: Any,
        target: Any,
    ) -> tuple[Any, dict[str, Any]]:
        diff = predicted - target
        sq = diff * diff
        if self.reduction == "mean":
            loss = sq.mean()
        elif self.reduction == "sum":
            loss = sq.sum()
        else:
            loss = sq
        return loss, {"mse": loss}

NLL

Bases: Operator

Negative log-likelihood for distribution-shaped predictions.

Contract: predicted must expose .log_prob(target) returning a per-sample log-probability array. The loss reduces over the leading batch axis with reduction.

Typical wrappers — your model emits a distribution-shaped output:

• For amortized inference: a pyrox / distrax / numpyro distribution whose .log_prob is differentiable through the model parameters.
• For Gaussian regression: a tiny head operator that wraps (mean, log_var) model outputs into a Normal-like object.

Parameters:

Name	Type	Description	Default
reduction	str	"mean" (default), "sum", or "none".	'mean'

Returns from __call__: (loss, {"nll": loss}).

Source code in packages/pipekit-train/src/pipekit_train/loss.py

class NLL(Operator):
    """Negative log-likelihood for distribution-shaped predictions.

    Contract: ``predicted`` must expose ``.log_prob(target)`` returning
    a per-sample log-probability array. The loss reduces over the
    leading batch axis with ``reduction``.

    Typical wrappers — your model emits a distribution-shaped output:

    - For amortized inference: a ``pyrox`` / ``distrax`` / ``numpyro``
      distribution whose ``.log_prob`` is differentiable through the
      model parameters.
    - For Gaussian regression: a tiny head operator that wraps
      ``(mean, log_var)`` model outputs into a ``Normal``-like object.

    Args:
        reduction: ``"mean"`` (default), ``"sum"``, or ``"none"``.

    Returns from ``__call__``:
        ``(loss, {"nll": loss})``.
    """

    def __init__(self, reduction: str = "mean") -> None:
        if reduction not in {"mean", "sum", "none"}:
            raise ValueError(
                f"NLL.reduction must be one of 'mean', 'sum', 'none'; "
                f"got {reduction!r}."
            )
        self.reduction = reduction

    def _apply(
        self,
        predicted: Any,
        target: Any,
    ) -> tuple[Any, dict[str, Any]]:
        if not hasattr(predicted, "log_prob"):
            raise TypeError(
                "NLL expects `predicted` to expose a `.log_prob(target)` "
                "method (a distribution-shaped output). Got "
                f"{type(predicted).__name__}. Wrap your model output in a "
                "density head, or use MSE for plain-array regression."
            )
        log_probs = predicted.log_prob(target)
        if self.reduction == "mean":
            loss = -log_probs.mean()
        elif self.reduction == "sum":
            loss = -log_probs.sum()
        else:
            loss = -log_probs
        return loss, {"nll": loss}

Callbacks

pipekit_train.callbacks

Callback protocol + bundled callbacks.

The five-hook callback surface and three bundled callbacks (EarlyStopping, Checkpoint, LogToExperiment) per docs/design/api/components.md.

The Callback Protocol is intentionally not @runtime_checkable: structural isinstance would reject partial implementations (a metric-logger that only cares about on_step_end), which contradicts the design's optional-hook model. Adapters dispatch via getattr(cb, hook, None). See ADR D9.

Callback

Bases: Protocol

Per-step / per-epoch / per-eval hooks.

Each of the five hooks is optional at runtime. Adapters dispatch via getattr(cb, hook_name, None) so a callback that only cares about one hook (say, on_step_end for a metric logger) implements just that method.

This Protocol is intentionally not @runtime_checkable: structural isinstance would require every callback to implement all five methods, which contradicts the partial-implementation model documented in docs/design/api/components.md. The Protocol exists as the type-checker contract describing the full hook surface; the runtime contract is "ducks with the methods you actually use".

See ADR D9 in docs/design/decisions.md.

Source code in packages/pipekit-train/src/pipekit_train/callbacks.py

class Callback(Protocol):
    """Per-step / per-epoch / per-eval hooks.

    Each of the five hooks is **optional** at runtime. Adapters
    dispatch via ``getattr(cb, hook_name, None)`` so a callback that
    only cares about one hook (say, ``on_step_end`` for a metric
    logger) implements just that method.

    This Protocol is intentionally **not** ``@runtime_checkable``:
    structural ``isinstance`` would require every callback to
    implement all five methods, which contradicts the
    partial-implementation model documented in
    ``docs/design/api/components.md``. The Protocol exists as the
    type-checker contract describing the *full* hook surface; the
    runtime contract is "ducks with the methods you actually use".

    See ADR D9 in ``docs/design/decisions.md``.
    """

    def on_train_begin(self, loop: Any, state: CarryState) -> None: ...

    def on_step_end(
        self,
        loop: Any,
        state: CarryState,
        metrics: dict[str, float],
    ) -> None: ...

    def on_epoch_end(
        self,
        loop: Any,
        state: CarryState,
        metrics: dict[str, float],
    ) -> None: ...

    def on_eval_end(
        self,
        loop: Any,
        state: CarryState,
        eval_metrics: dict[str, float],
    ) -> None: ...

    def on_train_end(self, loop: Any, state: CarryState) -> None: ...

Checkpoint

Bases: Operator

Save (model_op, optimizer_state) every N steps; final to a registry.

On each on_step_end, when state.step is a multiple of every_n_steps, the adapter is expected to save its native checkpoint into save_dir. (The actual save call lives in the adapter; this callback just signals when to fire.)

When registry is set, on_train_end stores the final model operator in the registry by content hash, and the URI is recorded on last_uri / last_hash for the loop to fold into its TrainingArtifact.

Parameters:

Name	Type	Description	Default
every_n_steps	int	Cadence for on-disk snapshots.	1000
keep_last	int	Number of recent checkpoints to retain on disk. Honored by the adapter's native checkpoint manager (Orbax CheckpointManager.max_to_keep, Lightning ModelCheckpoint.save_top_k).	3
save_dir	str	Local directory or fsspec URI.	'./ckpts'
registry	Any | None	Optional anything satisfying pipekit_experiment.ModelRegistry. When set, the final checkpoint is also stored by content hash and the URI is wired back into TrainingArtifact.model_registry_uri / .trained_model_hash.	None

Attributes (runtime state): last_uri: Set after on_train_end if a registry was provided. last_hash: Set after on_train_end if a registry was provided.

Source code in packages/pipekit-train/src/pipekit_train/callbacks.py

class Checkpoint(Operator):
    """Save (model_op, optimizer_state) every N steps; final to a registry.

    On each ``on_step_end``, when ``state.step`` is a multiple of
    ``every_n_steps``, the adapter is expected to save its native
    checkpoint into ``save_dir``. (The actual save call lives in the
    adapter; this callback just signals when to fire.)

    When ``registry`` is set, ``on_train_end`` stores the final model
    operator in the registry by content hash, and the URI is recorded
    on ``last_uri`` / ``last_hash`` for the loop to fold into its
    `TrainingArtifact`.

    Args:
        every_n_steps: Cadence for on-disk snapshots.
        keep_last: Number of recent checkpoints to retain on disk.
            Honored by the adapter's native checkpoint manager
            (Orbax `CheckpointManager.max_to_keep`, Lightning
            `ModelCheckpoint.save_top_k`).
        save_dir: Local directory or fsspec URI.
        registry: Optional anything satisfying
            ``pipekit_experiment.ModelRegistry``. When set, the final
            checkpoint is also stored by content hash and the URI is
            wired back into `TrainingArtifact.model_registry_uri` /
            `.trained_model_hash`.

    Attributes (runtime state):
        last_uri: Set after ``on_train_end`` if a registry was provided.
        last_hash: Set after ``on_train_end`` if a registry was provided.
    """

    __config_exclude__: ClassVar[tuple[str, ...]] = ("last_uri", "last_hash")
    # `registry` is a runtime object that doesn't round-trip cleanly.
    # We surface it as a constructor argument but mark the operator
    # forbid_in_yaml so YAML loaders refuse to round-trip a Checkpoint
    # that carries a non-serialisable registry handle.
    forbid_in_yaml: ClassVar[bool] = True

    def __init__(
        self,
        every_n_steps: int = 1000,
        keep_last: int = 3,
        save_dir: str = "./ckpts",
        registry: Any | None = None,
    ) -> None:
        if every_n_steps < 1:
            raise ValueError(
                f"Checkpoint.every_n_steps must be >= 1; got {every_n_steps}."
            )
        if keep_last < 1:
            raise ValueError(f"Checkpoint.keep_last must be >= 1; got {keep_last}.")
        self.every_n_steps = int(every_n_steps)
        self.keep_last = int(keep_last)
        self.save_dir = str(save_dir)
        self.registry = registry
        self.last_uri: str | None = None
        self.last_hash: str | None = None

    def should_snapshot(self, step: int) -> bool:
        """Pure predicate the adapter can call from on_step_end."""
        return step > 0 and step % self.every_n_steps == 0

    def on_train_end(self, loop: Any, state: Any) -> None:
        """Persist the final model to the registry, if one is configured."""
        del loop
        if self.registry is None:
            return
        model = getattr(state, "model", None)
        if model is None:
            return
        # `registry.store` is the pipekit_experiment.ModelRegistry protocol.
        # We don't import pipekit_experiment to keep this dependency
        # optional (`pipekit-train[experiment]`); duck-type instead.
        if not hasattr(self.registry, "store"):
            raise TypeError(
                "Checkpoint.registry must satisfy pipekit_experiment."
                "ModelRegistry (missing `.store`). Got "
                f"{type(self.registry).__name__}."
            )
        # If the model exposes `serialize_weights() -> bytes` (e.g.
        # `pipekit_jax.JaxModelOp`), pass the weight blob through so
        # the registry can store it alongside the operator state.
        # That's the round-trip path that closes the v0.1 reproducibility
        # caveat documented in `docs/design/boundaries.md §13.1` —
        # duck-typed so pipekit-train doesn't import pipekit-jax.
        serialize = getattr(model, "serialize_weights", None)
        if callable(serialize):
            h = self.registry.store(model, weights=serialize())
        else:
            h = self.registry.store(model)
        self.last_hash = h
        self.last_uri = f"registry://{h}"

    def _apply(self, *args: Any, **kwargs: Any) -> Any:
        raise TypeError(
            "Checkpoint is a callback, not a pipeline operator. "
            "Pass it via TrainingLoop(callbacks=...)."
        )

on_train_end(loop, state)

Persist the final model to the registry, if one is configured.

Source code in packages/pipekit-train/src/pipekit_train/callbacks.py

def on_train_end(self, loop: Any, state: Any) -> None:
    """Persist the final model to the registry, if one is configured."""
    del loop
    if self.registry is None:
        return
    model = getattr(state, "model", None)
    if model is None:
        return
    # `registry.store` is the pipekit_experiment.ModelRegistry protocol.
    # We don't import pipekit_experiment to keep this dependency
    # optional (`pipekit-train[experiment]`); duck-type instead.
    if not hasattr(self.registry, "store"):
        raise TypeError(
            "Checkpoint.registry must satisfy pipekit_experiment."
            "ModelRegistry (missing `.store`). Got "
            f"{type(self.registry).__name__}."
        )
    # If the model exposes `serialize_weights() -> bytes` (e.g.
    # `pipekit_jax.JaxModelOp`), pass the weight blob through so
    # the registry can store it alongside the operator state.
    # That's the round-trip path that closes the v0.1 reproducibility
    # caveat documented in `docs/design/boundaries.md §13.1` —
    # duck-typed so pipekit-train doesn't import pipekit-jax.
    serialize = getattr(model, "serialize_weights", None)
    if callable(serialize):
        h = self.registry.store(model, weights=serialize())
    else:
        h = self.registry.store(model)
    self.last_hash = h
    self.last_uri = f"registry://{h}"

should_snapshot(step)

Pure predicate the adapter can call from on_step_end.

Source code in packages/pipekit-train/src/pipekit_train/callbacks.py

def should_snapshot(self, step: int) -> bool:
    """Pure predicate the adapter can call from on_step_end."""
    return step > 0 and step % self.every_n_steps == 0

EarlyStopping

Bases: Operator

Signal stop when a monitored metric hasn't improved for N evals.

The callback only triggers on on_eval_end. It updates internal "best so far" state and sets should_stop to True once patience consecutive eval rounds pass without a meaningful improvement (min_delta).

Parameters:

Name	Type	Description	Default
metric	str	Metric name as it appears in the on_eval_end dict.	required
patience	int	Number of consecutive evals without improvement before signalling stop.	10
mode	Literal['min', 'max']	"min" for losses, "max" for accuracies.	'min'
min_delta	float	Minimum change to count as improvement.	0.0

Attributes (runtime state, not part of the operator config): should_stop: Read by the adapter's outer loop on each on_eval_end to decide whether to break early. best: Current best value of the monitored metric. wait: Consecutive eval rounds since last improvement.

Source code in packages/pipekit-train/src/pipekit_train/callbacks.py

class EarlyStopping(Operator):
    """Signal stop when a monitored metric hasn't improved for N evals.

    The callback only triggers on ``on_eval_end``. It updates internal
    "best so far" state and sets ``should_stop`` to True once
    ``patience`` consecutive eval rounds pass without a meaningful
    improvement (``min_delta``).

    Args:
        metric: Metric name as it appears in the on_eval_end dict.
        patience: Number of consecutive evals without improvement
            before signalling stop.
        mode: ``"min"`` for losses, ``"max"`` for accuracies.
        min_delta: Minimum change to count as improvement.

    Attributes (runtime state, not part of the operator config):
        should_stop: Read by the adapter's outer loop on each
            on_eval_end to decide whether to break early.
        best: Current best value of the monitored metric.
        wait: Consecutive eval rounds since last improvement.
    """

    # Runtime state lives on the operator but doesn't participate in
    # get_config (the dataclass-style discipline; tested below).
    __config_exclude__: ClassVar[tuple[str, ...]] = ("should_stop", "best", "wait")

    def __init__(
        self,
        metric: str,
        patience: int = 10,
        mode: Literal["min", "max"] = "min",
        min_delta: float = 0.0,
    ) -> None:
        if mode not in {"min", "max"}:
            raise ValueError(
                f"EarlyStopping.mode must be 'min' or 'max'; got {mode!r}."
            )
        if patience < 1:
            raise ValueError(f"EarlyStopping.patience must be >= 1; got {patience}.")
        if min_delta < 0:
            raise ValueError(
                f"EarlyStopping.min_delta must be non-negative; got {min_delta}."
            )
        self.metric = metric
        self.patience = int(patience)
        self.mode = mode
        self.min_delta = float(min_delta)
        # Runtime state. These are deliberately not in get_config.
        self.should_stop: bool = False
        self.best: float | None = None
        self.wait: int = 0

    def _is_improvement(self, value: float) -> bool:
        if self.best is None:
            return True
        if self.mode == "min":
            return value < self.best - self.min_delta
        return value > self.best + self.min_delta

    def on_eval_end(
        self,
        loop: Any,
        state: Any,
        eval_metrics: dict[str, float],
    ) -> None:
        del loop, state
        if self.metric not in eval_metrics:
            return  # silently skip; design choice — adapter need not error
        value = float(eval_metrics[self.metric])
        if self._is_improvement(value):
            self.best = value
            self.wait = 0
        else:
            self.wait += 1
            if self.wait >= self.patience:
                self.should_stop = True

    def reset(self) -> None:
        """Clear runtime state — useful for re-running the same loop."""
        self.should_stop = False
        self.best = None
        self.wait = 0

    def _apply(self, *args: Any, **kwargs: Any) -> Any:
        raise TypeError(
            "EarlyStopping is a callback, not a pipeline operator. "
            "Pass it via TrainingLoop(callbacks=...)."
        )

reset()

Clear runtime state — useful for re-running the same loop.

Source code in packages/pipekit-train/src/pipekit_train/callbacks.py

def reset(self) -> None:
    """Clear runtime state — useful for re-running the same loop."""
    self.should_stop = False
    self.best = None
    self.wait = 0

LogToExperiment

Bases: Operator

Forward metrics / params / artifacts to a pipekit-experiment tracker.

On on_train_begin, calls tracker.start_run with the loop's config (plus any user-supplied extra config). On on_step_end / on_eval_end, forwards metrics to tracker.log_metrics. On on_train_end, calls tracker.end_run.

Parameters:

Name	Type	Description	Default
tracker	Any	Anything satisfying pipekit_experiment.ExperimentTracker.	required
run_name	str	Name passed to tracker.start_run.	required
config	dict[str, Any] | None	Extra config dict merged into the run start. The tracker also sees the TrainingLoop's get_config().	None

Attributes (runtime state): run: The Run object returned by tracker.start_run, available after on_train_begin.

Source code in packages/pipekit-train/src/pipekit_train/callbacks.py

class LogToExperiment(Operator):
    """Forward metrics / params / artifacts to a pipekit-experiment tracker.

    On ``on_train_begin``, calls ``tracker.start_run`` with the loop's
    config (plus any user-supplied extra config). On ``on_step_end`` /
    ``on_eval_end``, forwards metrics to ``tracker.log_metrics``. On
    ``on_train_end``, calls ``tracker.end_run``.

    Args:
        tracker: Anything satisfying
            ``pipekit_experiment.ExperimentTracker``.
        run_name: Name passed to ``tracker.start_run``.
        config: Extra config dict merged into the run start. The
            tracker also sees the TrainingLoop's ``get_config()``.

    Attributes (runtime state):
        run: The Run object returned by ``tracker.start_run``,
            available after ``on_train_begin``.
    """

    __config_exclude__: ClassVar[tuple[str, ...]] = ("run",)
    forbid_in_yaml: ClassVar[bool] = True

    def __init__(
        self,
        tracker: Any,
        run_name: str,
        config: dict[str, Any] | None = None,
    ) -> None:
        for method in ("start_run", "log_metrics", "end_run"):
            if not hasattr(tracker, method):
                raise TypeError(
                    "LogToExperiment.tracker must satisfy pipekit_experiment."
                    f"ExperimentTracker (missing `.{method}`). Got "
                    f"{type(tracker).__name__}."
                )
        self.tracker = tracker
        self.run_name = run_name
        self.config = dict(config) if config is not None else None
        self.run: Any = None

    def on_train_begin(self, loop: Any, state: Any) -> None:
        del state
        run_config: dict[str, Any] = {}
        if loop is not None and hasattr(loop, "get_config"):
            run_config.update(loop.get_config())
        if self.config is not None:
            run_config.update(self.config)
        self.run = self.tracker.start_run(name=self.run_name, config=run_config)

    def on_step_end(
        self,
        loop: Any,
        state: Any,
        metrics: dict[str, float],
    ) -> None:
        del loop
        if self.run is None or not metrics:
            return
        step = getattr(state, "step", 0)
        self.tracker.log_metrics(self.run, dict(metrics), step=step)

    def on_eval_end(
        self,
        loop: Any,
        state: Any,
        eval_metrics: dict[str, float],
    ) -> None:
        del loop
        if self.run is None or not eval_metrics:
            return
        step = getattr(state, "step", 0)
        tagged = {f"eval/{k}": v for k, v in eval_metrics.items()}
        self.tracker.log_metrics(self.run, tagged, step=step)

    def on_train_end(self, loop: Any, state: Any) -> None:
        del loop, state
        if self.run is None:
            return
        self.tracker.end_run(self.run, status="completed")

    def _apply(self, *args: Any, **kwargs: Any) -> Any:
        raise TypeError(
            "LogToExperiment is a callback, not a pipeline operator. "
            "Pass it via TrainingLoop(callbacks=...)."
        )

Metric writers

pipekit_train.writer

Metric writer protocol + default JSON Lines implementation.

MetricWriter is the carrier-agnostic seam where per-step metrics land when no ExperimentTracker is attached. JSONLWriter is the default; W&B, TensorBoard, and other writers are user-supplied per docs/design/examples/integration.md.

JSONLWriter dataclass

Append-only JSON Lines metric writer.

Each write appends one line of the form::

{"step": int, "metrics": {name: float, ...}, "timestamp": iso8601}

The file is opened lazily on the first write. Records are written into Python's text-IO buffer immediately and that buffer is flushed to the OS every flush_every writes; flushing only pushes data out of user space, it does not guarantee durability against a kernel crash or power loss. Set fsync_on_flush=True to additionally call os.fsync on each flush — slower, but the data is on stable storage when it returns.

Parameters:

Name	Type	Description	Default
path	str | Path	Output file path. Parent directories are created on first write if missing.	required
flush_every	int	Number of writes between explicit flush() calls. Higher = fewer syscalls; lower = less data sitting in Python's buffer if the process is killed.	100
fsync_on_flush	bool	When True, os.fsync is invoked on each flush in addition to the buffer flush — provides durability against kernel/power loss at a cost.	False

See docs/design/api/components.md.

Source code in packages/pipekit-train/src/pipekit_train/writer.py

@dataclass
class JSONLWriter:
    """Append-only JSON Lines metric writer.

    Each ``write`` appends one line of the form::

        {"step": int, "metrics": {name: float, ...}, "timestamp": iso8601}

    The file is opened lazily on the first ``write``. Records are
    written into Python's text-IO buffer immediately and that buffer
    is flushed to the OS every ``flush_every`` writes; flushing only
    pushes data out of user space, it does **not** guarantee
    durability against a kernel crash or power loss. Set
    ``fsync_on_flush=True`` to additionally call ``os.fsync`` on each
    flush — slower, but the data is on stable storage when it returns.

    Args:
        path: Output file path. Parent directories are created on
            first write if missing.
        flush_every: Number of writes between explicit ``flush()``
            calls. Higher = fewer syscalls; lower = less data sitting
            in Python's buffer if the process is killed.
        fsync_on_flush: When True, ``os.fsync`` is invoked on each
            flush in addition to the buffer flush — provides
            durability against kernel/power loss at a cost.

    See ``docs/design/api/components.md``.
    """

    path: str | Path
    flush_every: int = 100
    fsync_on_flush: bool = False

    _file: Any = field(default=None, init=False, repr=False, compare=False)
    _writes_since_flush: int = field(default=0, init=False, repr=False, compare=False)

    def write(self, step: int, metrics: dict[str, float]) -> None:
        """Append one record."""
        if self._file is None:
            path = Path(self.path)
            path.parent.mkdir(parents=True, exist_ok=True)
            self._file = path.open("a", encoding="utf-8")
        record = {
            "step": int(step),
            "metrics": {k: float(v) for k, v in metrics.items()},
            "timestamp": datetime.now(UTC).isoformat(),
        }
        self._file.write(json.dumps(record) + "\n")
        self._writes_since_flush += 1
        if self._writes_since_flush >= self.flush_every:
            self._flush()
            self._writes_since_flush = 0

    def close(self) -> None:
        """Close the underlying file."""
        if self._file is not None:
            self._flush()
            self._file.close()
            self._file = None
            self._writes_since_flush = 0

    def _flush(self) -> None:
        """Flush the text-IO buffer and optionally fsync."""
        assert self._file is not None
        self._file.flush()
        if self.fsync_on_flush:
            os.fsync(self._file.fileno())

close()

Close the underlying file.

Source code in packages/pipekit-train/src/pipekit_train/writer.py

def close(self) -> None:
    """Close the underlying file."""
    if self._file is not None:
        self._flush()
        self._file.close()
        self._file = None
        self._writes_since_flush = 0

write(step, metrics)

Append one record.

Source code in packages/pipekit-train/src/pipekit_train/writer.py

def write(self, step: int, metrics: dict[str, float]) -> None:
    """Append one record."""
    if self._file is None:
        path = Path(self.path)
        path.parent.mkdir(parents=True, exist_ok=True)
        self._file = path.open("a", encoding="utf-8")
    record = {
        "step": int(step),
        "metrics": {k: float(v) for k, v in metrics.items()},
        "timestamp": datetime.now(UTC).isoformat(),
    }
    self._file.write(json.dumps(record) + "\n")
    self._writes_since_flush += 1
    if self._writes_since_flush >= self.flush_every:
        self._flush()
        self._writes_since_flush = 0

MetricWriter

Bases: Protocol

Where per-step metrics go when no tracker is attached.

Implementations: JSONL (default), in-memory list, W&B, TensorBoard. The trainer calls write per logged step.

See docs/design/api/components.md.

Source code in packages/pipekit-train/src/pipekit_train/writer.py

@runtime_checkable
class MetricWriter(Protocol):
    """Where per-step metrics go when no tracker is attached.

    Implementations: JSONL (default), in-memory list, W&B,
    TensorBoard. The trainer calls ``write`` per logged step.

    See ``docs/design/api/components.md``.
    """

    def write(self, step: int, metrics: dict[str, float]) -> None: ...

    def close(self) -> None: ...

Training loop

pipekit_train.loop

Layer 3 — TrainingLoop.

The headline composable. TrainingLoop is a pipekit.StatefulOperator carrying (model, opt_state, step, epoch, metrics, rng); running it produces (trained_model_op, artifact).

This module also ships the supporting TrainerCarryState (subclass of pipekit.CarryState) and ValidationStep (an Operator that evaluates a trained model against a held-out dataset).

See docs/design/api/models.md.

TrainerCarryState

Bases: CarryState

Carry-state for TrainingLoop.

The optimizer state and RNG are backend-specific opaque types (an optax.OptState and a jax.Array PRNG key for the Equinox adapter; a torch.optim.Optimizer.state_dict() + tuple for Lightning). The pipekit-train side carries them as Any and delegates serialisation to the adapter's native checkpoint machinery.

Attributes:

Name	Type	Description
model		Current model operator (possibly partially trained).
opt_state		Backend-specific optimizer state.
step		Global step counter.
epoch		Derived from step // steps_per_epoch (0 if not set).
metrics		Latest per-step metric snapshot (cleared on epoch end).
rng		Backend-specific PRNG state.

Source code in packages/pipekit-train/src/pipekit_train/loop.py

class TrainerCarryState(CarryState):
    """Carry-state for `TrainingLoop`.

    The optimizer state and RNG are backend-specific opaque types
    (an `optax.OptState` and a `jax.Array` PRNG key for the Equinox
    adapter; a `torch.optim.Optimizer.state_dict()` + tuple for
    Lightning). The pipekit-train side carries them as `Any` and
    delegates serialisation to the adapter's native checkpoint
    machinery.

    Attributes:
        model: Current model operator (possibly partially trained).
        opt_state: Backend-specific optimizer state.
        step: Global step counter.
        epoch: Derived from step // steps_per_epoch (0 if not set).
        metrics: Latest per-step metric snapshot (cleared on epoch end).
        rng: Backend-specific PRNG state.
    """

    def __init__(
        self,
        model: Operator,
        opt_state: Any,
        step: int = 0,
        epoch: int = 0,
        metrics: dict[str, float] | None = None,
        rng: Any = None,
    ) -> None:
        self.model = model
        self.opt_state = opt_state
        self.step = int(step)
        self.epoch = int(epoch)
        self.metrics = dict(metrics) if metrics is not None else {}
        self.rng = rng

    def to_dict(self) -> dict[str, Any]:
        """JSON-safe snapshot.

        Backend-specific fields (``model``, ``opt_state``, ``rng``)
        are stored as repr strings — the adapter's native checkpoint
        is the round-trip path for those.
        """
        return {
            "model": repr(self.model),
            "opt_state": repr(self.opt_state),
            "step": self.step,
            "epoch": self.epoch,
            "metrics": dict(self.metrics),
            "rng": repr(self.rng),
        }

to_dict()

JSON-safe snapshot.

Backend-specific fields (model, opt_state, rng) are stored as repr strings — the adapter's native checkpoint is the round-trip path for those.

Source code in packages/pipekit-train/src/pipekit_train/loop.py

def to_dict(self) -> dict[str, Any]:
    """JSON-safe snapshot.

    Backend-specific fields (``model``, ``opt_state``, ``rng``)
    are stored as repr strings — the adapter's native checkpoint
    is the round-trip path for those.
    """
    return {
        "model": repr(self.model),
        "opt_state": repr(self.opt_state),
        "step": self.step,
        "epoch": self.epoch,
        "metrics": dict(self.metrics),
        "rng": repr(self.rng),
    }

TrainingLoop

Bases: StatefulOperator

Train a model_op on a dataset, producing a trained model_op.

The headline composable. Delegates the actual training to a backend adapter (pipekit_train.adapters.<backend>.run). The adapter sees the full loop config and is responsible for the per-step gradient update and the per-eval validation pass. The TrainingLoop itself owns the artifact assembly and the callback threading.

Parameters:

Name	Type	Description	Default
model_op	Operator	Untrained model Operator. Backend-specific (an equinox/lightning/keras module wrapped in a pipekit.Operator).	required
dataset	TrainingDataset	A TrainingDataset.	required
val_dataset	TrainingDataset | None	Optional held-out TrainingDataset for eval.	None
loss	Loss | None	A Loss. Ignored if task is supplied.	None
task	Any | None	Optional backend-native task interface. Lets power users bypass the auto-synthesised Loss path.	None
optimizer_config	dict[str, Any] | None	Flat dict, e.g. {"name": "adamw", "lr": 3e-4, "weight_decay": 0.01}. Translated by the adapter into the backend's optimizer.	None
max_steps	int	Total optimisation steps. (Per-step is the unit; see ADR D10.)	10000
steps_per_epoch	int | None	Optional. Controls when on_epoch_end fires.	None
batch_size	int	Per-step minibatch size.	32
backend	Literal['equinox', 'lightning', 'keras', 'numpyro-svi', 'numpyro-mcmc', 'blackjax']	One of "equinox", "lightning", "keras", "numpyro-svi", "numpyro-mcmc", "blackjax".	'equinox'
callbacks	tuple[Any, ...]	Tuple of Callback instances.	()
metric_writer	MetricWriter | None	Optional MetricWriter.	None
eval_every_n_steps	int	Validation cadence.	500
log_every_n_steps	int	Metric logging cadence.	100
checkpoint_dir	str | None	Optional. When set, the adapter wires up its native checkpoint machinery here.	None
sharding	Any | None	Optional backend-specific sharding configuration for multi-device / multi-host training. For the Equinox backend this is a pipekit_train.adapters.equinox.ShardingSpec; None keeps single-device behaviour. Typed loosely here so the carrier-agnostic core never imports JAX.	None
seed	int	Master PRNG seed.	0

Raises (at run-time): ValueError: if neither loss nor task is supplied. ValueError: if backend isn't recognised.

Source code in packages/pipekit-train/src/pipekit_train/loop.py

class TrainingLoop(StatefulOperator):
    """Train a ``model_op`` on a ``dataset``, producing a trained model_op.

    The headline composable. Delegates the actual training to a
    backend adapter (``pipekit_train.adapters.<backend>.run``). The
    adapter sees the full loop config and is responsible for the
    per-step gradient update and the per-eval validation pass. The
    `TrainingLoop` itself owns the artifact assembly and the callback
    threading.

    Args:
        model_op: Untrained model Operator. Backend-specific (an
            equinox/lightning/keras module wrapped in a
            pipekit.Operator).
        dataset: A `TrainingDataset`.
        val_dataset: Optional held-out `TrainingDataset` for eval.
        loss: A `Loss`. Ignored if ``task`` is supplied.
        task: Optional backend-native task interface. Lets power users
            bypass the auto-synthesised `Loss` path.
        optimizer_config: Flat dict, e.g.
            ``{"name": "adamw", "lr": 3e-4, "weight_decay": 0.01}``.
            Translated by the adapter into the backend's optimizer.
        max_steps: Total optimisation steps. (Per-step is the unit;
            see ADR D10.)
        steps_per_epoch: Optional. Controls when on_epoch_end fires.
        batch_size: Per-step minibatch size.
        backend: One of ``"equinox"``, ``"lightning"``, ``"keras"``,
            ``"numpyro-svi"``, ``"numpyro-mcmc"``, ``"blackjax"``.
        callbacks: Tuple of `Callback` instances.
        metric_writer: Optional `MetricWriter`.
        eval_every_n_steps: Validation cadence.
        log_every_n_steps: Metric logging cadence.
        checkpoint_dir: Optional. When set, the adapter wires up its
            native checkpoint machinery here.
        sharding: Optional backend-specific sharding configuration for
            multi-device / multi-host training. For the Equinox backend
            this is a ``pipekit_train.adapters.equinox.ShardingSpec``;
            ``None`` keeps single-device behaviour. Typed loosely here so
            the carrier-agnostic core never imports JAX.
        seed: Master PRNG seed.

    Raises (at run-time):
        ValueError: if neither ``loss`` nor ``task`` is supplied.
        ValueError: if ``backend`` isn't recognised.
    """

    __config_mixin_auto__ = False
    # The training loop carries closures (model_op, dataset, callbacks)
    # that often don't round-trip cleanly; mark accordingly.
    forbid_in_yaml: ClassVar[bool] = True

    def __init__(
        self,
        model_op: Operator,
        dataset: TrainingDataset,
        val_dataset: TrainingDataset | None = None,
        loss: Loss | None = None,
        task: Any | None = None,
        optimizer_config: dict[str, Any] | None = None,
        max_steps: int = 10_000,
        steps_per_epoch: int | None = None,
        batch_size: int = 32,
        backend: Literal[
            "equinox",
            "lightning",
            "keras",
            "numpyro-svi",
            "numpyro-mcmc",
            "blackjax",
        ] = "equinox",
        callbacks: tuple[Any, ...] = (),
        metric_writer: MetricWriter | None = None,
        eval_every_n_steps: int = 500,
        log_every_n_steps: int = 100,
        checkpoint_dir: str | None = None,
        sharding: Any | None = None,
        seed: int = 0,
    ) -> None:
        if not isinstance(model_op, Operator):
            raise TypeError(
                "TrainingLoop.model_op must be an Operator; got "
                f"{type(model_op).__name__}."
            )
        if not isinstance(dataset, TrainingDataset):
            raise TypeError(
                "TrainingLoop.dataset must be a TrainingDataset; got "
                f"{type(dataset).__name__}."
            )
        if val_dataset is not None and not isinstance(val_dataset, TrainingDataset):
            raise TypeError(
                "TrainingLoop.val_dataset must be a TrainingDataset or None; "
                f"got {type(val_dataset).__name__}."
            )
        if backend not in _BACKEND_MODULES:
            raise ValueError(
                f"TrainingLoop.backend must be one of "
                f"{sorted(_BACKEND_MODULES)}; got {backend!r}."
            )
        if max_steps < 1:
            raise ValueError(f"TrainingLoop.max_steps must be >= 1; got {max_steps}.")
        if batch_size < 1:
            raise ValueError(f"TrainingLoop.batch_size must be >= 1; got {batch_size}.")
        self.model_op = model_op
        self.dataset = dataset
        self.val_dataset = val_dataset
        self.loss = loss
        self.task = task
        self.optimizer_config = (
            dict(optimizer_config)
            if optimizer_config
            else {
                "name": "adamw",
                "lr": 3e-4,
            }
        )
        self.max_steps = int(max_steps)
        self.steps_per_epoch = steps_per_epoch
        self.batch_size = int(batch_size)
        self.backend = backend
        self.callbacks = tuple(callbacks)
        self.metric_writer = metric_writer
        self.eval_every_n_steps = int(eval_every_n_steps)
        self.log_every_n_steps = int(log_every_n_steps)
        self.checkpoint_dir = checkpoint_dir
        self.sharding = sharding
        self.seed = int(seed)

    # --- StatefulOperator hook ----------------------------------------

    def _apply(
        self,
        carrier: Any = None,
        state: TrainerCarryState | None = None,
    ) -> tuple[Operator, TrainerCarryState]:
        """Run training; return ``(trained_model_op, final_state)``.

        v0.1 scope: ``carrier`` is ignored (the loop's own ``model_op``
        is always used). ``state`` is also ignored (resume happens
        through the adapter's native checkpoint manager via
        ``checkpoint_dir``). Sequential composition with prior-stage
        model handoff is documented but unsupported until v0.2.
        """
        del carrier, state
        trained_model_op, backend_info = self._run_adapter()
        # Honour the adapter's actual final_step — the loop may have
        # terminated early (EarlyStopping) and self.max_steps would
        # otherwise inflate step/epoch in the returned carry-state.
        final_step = int(backend_info.get("final_step", self.max_steps))
        final_epoch = final_step // self.steps_per_epoch if self.steps_per_epoch else 0
        final_state = TrainerCarryState(
            model=trained_model_op,
            opt_state=None,
            step=final_step,
            epoch=final_epoch,
            metrics=dict(backend_info.get("final_metrics", {})),
            rng=None,
        )
        # The artifact build path uses _run_adapter's backend_info too.
        # Re-use the backend_info dict so .run() can pull it without
        # re-running adapter.run().
        self._last_backend_info = backend_info
        return trained_model_op, final_state

    # --- Convenience entry --------------------------------------------

    def run(self) -> tuple[Operator, Any]:
        """Run training end-to-end and return ``(trained_model_op, artifact)``.

        Equivalent to ``self._apply(None, None)`` followed by an
        artifact build. The artifact is a
        `pipekit_experiment.TrainingArtifact` when that package is
        installed; otherwise a local `_LocalArtifact` dataclass with
        the same field names.
        """
        trained_model_op, _ = self._apply()
        artifact = self._build_artifact(
            trained_model_op,
            getattr(self, "_last_backend_info", {}),
        )
        return trained_model_op, artifact

    # --- Adapter dispatch ---------------------------------------------

    def _run_adapter(self) -> tuple[Operator, dict[str, Any]]:
        """Resolve and call the configured backend adapter."""
        module_name = _BACKEND_MODULES[self.backend]
        module = import_module(module_name)
        return module.run(self)

    # --- Artifact assembly --------------------------------------------

    def _build_artifact(
        self,
        trained_model_op: Operator,
        backend_info: dict[str, Any],
    ) -> Any:
        """Assemble the `TrainingArtifact` (or `_LocalArtifact` fallback)."""
        # Pull tracker/registry references from callbacks (if any).
        tracker_run_id: str | None = None
        model_registry_uri: str = ""
        trained_model_hash: str = ""
        for cb in self.callbacks:
            run = getattr(cb, "run", None)
            if run is not None and hasattr(run, "__getitem__"):
                # LogToExperiment.run is a tracker-specific record.
                # We surface a stable string representation.
                tracker_run_id = repr(run)
            last_uri = getattr(cb, "last_uri", None)
            last_hash = getattr(cb, "last_hash", None)
            if last_uri:
                model_registry_uri = last_uri
            if last_hash:
                trained_model_hash = last_hash

        artifact_cls: Any
        try:
            artifact_cls = import_module(
                "pipekit_experiment.artifacts"
            ).TrainingArtifact
        except ImportError:
            artifact_cls = _LocalArtifact

        # `pipekit.dumps` already returns a JSON string; don't re-encode it.
        # Pretty-print by round-tripping through json.loads for readability.
        try:
            training_pipeline_yaml = json.dumps(
                json.loads(dumps(self)), indent=2, sort_keys=True
            )
        except Exception:
            # forbid_in_yaml means we may not round-trip cleanly;
            # surface a JSON repr of get_config() instead.
            training_pipeline_yaml = json.dumps(
                self.get_config(), indent=2, default=repr
            )

        return artifact_cls(
            training_pipeline_yaml=training_pipeline_yaml,
            dataset_hash=self.dataset.content_hash(),
            trained_model_hash=trained_model_hash,
            tracker_run_id=tracker_run_id,
            model_registry_uri=model_registry_uri,
            backend_info=backend_info,
            deps_lock=_capture_deps_lock(),
            metadata={},
        )

    # --- Config / repr -------------------------------------------------

    def get_config(self) -> dict[str, Any]:
        return {
            "model_op": type(self.model_op).__name__,
            "dataset": type(self.dataset).__name__,
            "dataset_hash": self.dataset.content_hash(),
            "val_dataset": (
                type(self.val_dataset).__name__
                if self.val_dataset is not None
                else None
            ),
            "loss": type(self.loss).__name__ if self.loss is not None else None,
            "task": type(self.task).__name__ if self.task is not None else None,
            "optimizer_config": dict(self.optimizer_config),
            "max_steps": self.max_steps,
            "steps_per_epoch": self.steps_per_epoch,
            "batch_size": self.batch_size,
            "backend": self.backend,
            "callbacks": [type(c).__name__ for c in self.callbacks],
            "eval_every_n_steps": self.eval_every_n_steps,
            "log_every_n_steps": self.log_every_n_steps,
            "checkpoint_dir": self.checkpoint_dir,
            "sharding": repr(self.sharding) if self.sharding is not None else None,
            "seed": self.seed,
        }

run()

Run training end-to-end and return (trained_model_op, artifact).

Equivalent to self._apply(None, None) followed by an artifact build. The artifact is a pipekit_experiment.TrainingArtifact when that package is installed; otherwise a local _LocalArtifact dataclass with the same field names.

Source code in packages/pipekit-train/src/pipekit_train/loop.py

def run(self) -> tuple[Operator, Any]:
    """Run training end-to-end and return ``(trained_model_op, artifact)``.

    Equivalent to ``self._apply(None, None)`` followed by an
    artifact build. The artifact is a
    `pipekit_experiment.TrainingArtifact` when that package is
    installed; otherwise a local `_LocalArtifact` dataclass with
    the same field names.
    """
    trained_model_op, _ = self._apply()
    artifact = self._build_artifact(
        trained_model_op,
        getattr(self, "_last_backend_info", {}),
    )
    return trained_model_op, artifact

ValidationStep

Bases: Operator

Compute validation metrics on a held-out dataset.

Iterates dataset, calls model_op on each input, applies each metric Operator to (predicted, target), and returns the mean-over-batches metric values.

Parameters:

Name	Type	Description	Default
model_op	Operator	The model Operator (trained or otherwise).	required
dataset	TrainingDataset	The held-out TrainingDataset.	required
metrics	tuple[Operator, ...]	Tuple of Operators. Each takes (predicted, target) and returns either a scalar or (scalar, aux_dict); the scalar is averaged. The aux dict from MSE etc. is also averaged per-key.	required

Returns from __call__: dict[str, float] with one entry per metric (plus any aux keys), each averaged over the dataset.

Source code in packages/pipekit-train/src/pipekit_train/loop.py

class ValidationStep(Operator):
    """Compute validation metrics on a held-out dataset.

    Iterates ``dataset``, calls ``model_op`` on each input, applies
    each metric `Operator` to ``(predicted, target)``, and returns
    the mean-over-batches metric values.

    Args:
        model_op: The model `Operator` (trained or otherwise).
        dataset: The held-out `TrainingDataset`.
        metrics: Tuple of `Operator`s. Each takes ``(predicted, target)``
            and returns either a scalar or ``(scalar, aux_dict)``; the
            scalar is averaged. The aux dict from `MSE` etc. is
            *also* averaged per-key.

    Returns from ``__call__``:
        ``dict[str, float]`` with one entry per metric (plus any aux
        keys), each averaged over the dataset.
    """

    __config_mixin_auto__ = False
    forbid_in_yaml: ClassVar[bool] = True

    def __init__(
        self,
        model_op: Operator,
        dataset: TrainingDataset,
        metrics: tuple[Operator, ...],
    ) -> None:
        if not isinstance(model_op, Operator):
            raise TypeError(
                "ValidationStep.model_op must be an Operator; got "
                f"{type(model_op).__name__}."
            )
        if not isinstance(dataset, TrainingDataset):
            raise TypeError(
                "ValidationStep.dataset must be a TrainingDataset; got "
                f"{type(dataset).__name__}."
            )
        if not metrics:
            raise ValueError("ValidationStep.metrics must be non-empty.")
        self.model_op = model_op
        self.dataset = dataset
        self.metrics = tuple(metrics)

    def _apply(self) -> dict[str, float]:
        sums: dict[str, float] = {}
        count = 0
        for x, y in self.dataset:
            predicted = self.model_op(x)
            for metric in self.metrics:
                result = metric(predicted, y)
                name = type(metric).__name__
                if isinstance(result, tuple):
                    scalar, aux = result
                    # Aggregate the primary scalar under the metric
                    # class name (in addition to the aux dict) so a
                    # metric that doesn't repeat its scalar in aux
                    # still has a recoverable headline value.
                    sums[name] = sums.get(name, 0.0) + float(scalar)
                    for k, v in aux.items():
                        sums[k] = sums.get(k, 0.0) + float(v)
                else:
                    sums[name] = sums.get(name, 0.0) + float(result)
            count += 1
        if count == 0:
            return {}
        return {k: v / count for k, v in sums.items()}

    def get_config(self) -> dict[str, Any]:
        return {
            "model_op": type(self.model_op).__name__,
            "dataset": type(self.dataset).__name__,
            "metrics": [type(m).__name__ for m in self.metrics],
        }

Hyperparameter sweeps

pipekit_train.sweep

Hyperparameter sweep orchestration.

Two operators:

• ParameterGrid — Cartesian-product iterator over named parameter axes. Yields {name: value, ...} dicts.
• HyperSweep — orchestrates N independent TrainingLoop runs by calling a user-supplied loop_template(params) -> TrainingLoop for each point in the search space, then returns the trial results ranked by a configured metric.

v0.2 ships grid + sequential execution. Optuna / Ray Tune sampler adapters land in v0.2.1; parallel execution lands in v0.3. The StatefulOperator wrapping means a sweep composes into Sequential just like a TrainingLoop (carry-state is the current best trial).

See issue #14 and design boundaries.md Q5.

HyperSweep

Bases: StatefulOperator

Orchestrate N independent TrainingLoop runs.

For each params dict yielded by search_space, builds a TrainingLoop via loop_template(params) and runs it. Captures the trained model and artifact per trial. Returns a list of SweepResult sorted by select_metric in select_mode order.

v0.2 runs trials sequentially. Parallel execution (multiprocessing) lands in v0.3. Optuna / Ray Tune samplers land in v0.2.1 — for now any iterable of param dicts works as search_space.

Parameters:

Name	Type	Description	Default
loop_template	Callable[[dict[str, Any]], TrainingLoop]	Callable (params: dict) -> TrainingLoop. Each trial's loop is built fresh; loop_template should be deterministic given params for reproducibility.	required
search_space	Iterable[dict[str, Any]]	A ParameterGrid or any iterable of param dicts. The latter accommodates Optuna trial generators in v0.2.1.	required
select_metric	str	Key to look up in artifact.backend_info["final_metrics"]. Default "loss" — match this to whatever your loss surfaces (e.g. "mse" for MSE).	'loss'
select_mode	Literal['min', 'max']	"min" for losses, "max" for accuracies.	'min'
n_trials	int | None	Optional cap on trials drawn from search_space. None runs the entire search space. Useful when search_space is an unbounded sampler.	None
sweep_id	str | None	Stable identifier for this sweep; default uuid4 hex. Surfaced on every SweepResult for tracker provenance.	None
on_trial	Callable[[SweepResult, SweepCarryState], None] | None	Optional callback (SweepResult, SweepCarryState) -> None invoked after each trial completes. Useful for progress logging or early-stop heuristics.	None

Returns from run(): list[SweepResult] sorted by select_metric. Best result is results[0].

Raises:

Type	Description
ValueError	if the search space is empty, or if a trial's artifact doesn't include select_metric.
TypeError	if loop_template(params) doesn't return a TrainingLoop.

Source code in packages/pipekit-train/src/pipekit_train/sweep.py

class HyperSweep(StatefulOperator):
    """Orchestrate N independent `TrainingLoop` runs.

    For each ``params`` dict yielded by ``search_space``, builds a
    `TrainingLoop` via ``loop_template(params)`` and runs it.
    Captures the trained model and artifact per trial. Returns a list
    of `SweepResult` sorted by ``select_metric`` in ``select_mode``
    order.

    v0.2 runs trials sequentially. Parallel execution (multiprocessing)
    lands in v0.3. Optuna / Ray Tune samplers land in v0.2.1 — for now
    any iterable of param dicts works as ``search_space``.

    Args:
        loop_template: Callable ``(params: dict) -> TrainingLoop``.
            Each trial's loop is built fresh; ``loop_template`` should
            be deterministic given ``params`` for reproducibility.
        search_space: A `ParameterGrid` or any iterable of param
            dicts. The latter accommodates Optuna trial generators in
            v0.2.1.
        select_metric: Key to look up in
            ``artifact.backend_info["final_metrics"]``. Default
            ``"loss"`` — match this to whatever your loss surfaces
            (e.g. ``"mse"`` for `MSE`).
        select_mode: ``"min"`` for losses, ``"max"`` for accuracies.
        n_trials: Optional cap on trials drawn from ``search_space``.
            ``None`` runs the entire search space. Useful when
            ``search_space`` is an unbounded sampler.
        sweep_id: Stable identifier for this sweep; default uuid4 hex.
            Surfaced on every `SweepResult` for tracker provenance.
        on_trial: Optional callback ``(SweepResult, SweepCarryState)
            -> None`` invoked after each trial completes. Useful for
            progress logging or early-stop heuristics.

    Returns from `run()`:
        ``list[SweepResult]`` sorted by ``select_metric``. Best result
        is ``results[0]``.

    Raises:
        ValueError: if the search space is empty, or if a trial's
            artifact doesn't include ``select_metric``.
        TypeError: if ``loop_template(params)`` doesn't return a
            `TrainingLoop`.
    """

    # The sweep itself isn't directly YAML-serialisable because the
    # loop_template is a user closure; mark accordingly.
    forbid_in_yaml: ClassVar[bool] = True
    __config_mixin_auto__ = False

    def __init__(
        self,
        loop_template: Callable[[dict[str, Any]], TrainingLoop],
        search_space: Iterable[dict[str, Any]],
        select_metric: str = "loss",
        select_mode: Literal["min", "max"] = "min",
        n_trials: int | None = None,
        sweep_id: str | None = None,
        on_trial: Callable[[SweepResult, SweepCarryState], None] | None = None,
    ) -> None:
        if not callable(loop_template):
            raise TypeError(
                "HyperSweep.loop_template must be callable; got "
                f"{type(loop_template).__name__}."
            )
        if select_mode not in {"min", "max"}:
            raise ValueError(
                f"HyperSweep.select_mode must be 'min' or 'max'; got {select_mode!r}."
            )
        if n_trials is not None and n_trials < 1:
            raise ValueError(
                f"HyperSweep.n_trials must be >= 1 or None; got {n_trials}."
            )
        self.loop_template = loop_template
        self.search_space = search_space
        self.select_metric = select_metric
        self.select_mode = select_mode
        self.n_trials = n_trials
        self.sweep_id = sweep_id or uuid.uuid4().hex
        self.on_trial = on_trial
        self.initial_state_fn = SweepCarryState

    # --- StatefulOperator hook -----------------------------------------

    def _apply(
        self,
        carrier: Any = None,
        state: SweepCarryState | None = None,
    ) -> tuple[list[SweepResult], SweepCarryState]:
        """Run all trials; return sorted results + final carry-state.

        ``carrier`` is ignored. ``state`` may be a `SweepCarryState`
        from a prior `_apply` call (used in `Sequential` composition);
        completed counts accumulate.
        """
        del carrier
        state = state or SweepCarryState()
        # Trial indices are GLOBAL across the cumulative trial stream
        # (state.completed at entry), not stage-local. That keeps
        # `SweepResult.trial_index` and `state.best_trial_index`
        # meaningful when a HyperSweep is composed in Sequential and
        # resumes from a non-zero `state.completed` (e.g. coarse →
        # fine sweep stages). Fresh sweeps (state.completed == 0)
        # see indices 0, 1, 2, … as before.
        starting_count = state.completed

        results: list[SweepResult] = []
        for local_index, params in enumerate(self._iter_trials()):
            global_index = starting_count + local_index
            loop = self.loop_template(dict(params))
            if not isinstance(loop, TrainingLoop):
                raise TypeError(
                    "HyperSweep.loop_template must return a "
                    f"TrainingLoop; got {type(loop).__name__}."
                )
            trained_model_op, artifact = loop.run()
            metric_value = self._extract_metric(artifact, global_index, params)

            result = SweepResult(
                params=dict(params),
                trained_model_op=trained_model_op,
                artifact=artifact,
                final_metric=metric_value,
                sweep_id=self.sweep_id,
                trial_index=global_index,
            )
            results.append(result)

            # Update carry-state.
            state.completed += 1
            if state.best_metric is None or self._is_better(
                metric_value, state.best_metric
            ):
                state.best_metric = metric_value
                state.best_trial_index = global_index

            if self.on_trial is not None:
                self.on_trial(result, state)

        sorted_results = sorted(
            results,
            key=lambda r: r.final_metric,
            reverse=self.select_mode == "max",
        )
        return sorted_results, state

    # --- run convenience -----------------------------------------------

    def run(self) -> list[SweepResult]:
        """Run all trials and return them sorted by ``select_metric``."""
        results, _ = self._apply()
        return results

    # --- helpers --------------------------------------------------------

    def _iter_trials(self) -> Iterator[dict[str, Any]]:
        seen = 0
        for params in iter(self.search_space):
            if not isinstance(params, dict):
                raise TypeError(
                    "HyperSweep.search_space must yield dicts; got "
                    f"{type(params).__name__} on trial {seen}."
                )
            yield params
            seen += 1
            if self.n_trials is not None and seen >= self.n_trials:
                return
        if seen == 0:
            raise ValueError("HyperSweep.search_space is empty — no trials to run.")

    def _extract_metric(
        self,
        artifact: Any,
        trial_index: int,
        params: dict[str, Any],
    ) -> float:
        backend_info = getattr(artifact, "backend_info", {}) or {}
        final_metrics = backend_info.get("final_metrics", {}) or {}
        if self.select_metric not in final_metrics:
            raise ValueError(
                f"HyperSweep: trial {trial_index} (params={params!r}) "
                f"artifact has no metric {self.select_metric!r} in "
                f"backend_info['final_metrics']. Available keys: "
                f"{sorted(final_metrics)}."
            )
        return float(final_metrics[self.select_metric])

    def _is_better(self, new: float, current_best: float) -> bool:
        if self.select_mode == "min":
            return new < current_best
        return new > current_best

    def get_config(self) -> dict[str, Any]:
        return {
            "select_metric": self.select_metric,
            "select_mode": self.select_mode,
            "n_trials": self.n_trials,
            "sweep_id": self.sweep_id,
            "search_space": type(self.search_space).__name__,
            "loop_template": getattr(
                self.loop_template, "__name__", repr(self.loop_template)
            ),
        }

run()

Run all trials and return them sorted by select_metric.

Source code in packages/pipekit-train/src/pipekit_train/sweep.py

def run(self) -> list[SweepResult]:
    """Run all trials and return them sorted by ``select_metric``."""
    results, _ = self._apply()
    return results

ParameterGrid

Bases: Operator

Cartesian product over named parameter axes.

Parameters:

Name	Type	Description	Default
grid	dict[str, list[Any]]	Mapping axis_name → list of values. Every axis must have at least one value. An empty mapping ({}) is also allowed and yields zero combinations (len(grid) == 0, list(grid) == []) — this departs from the mathematical Cartesian-product convention where ∏ over zero axes == 1, in favour of the pragmatic "no axes ⇒ no trials to run" reading that downstream HyperSweep consumers expect.	required

Iteration yields one {axis_name: value, ...} dict per cell. len(grid) is the product of axis lengths for non-empty grids; 0 for the empty grid per the convention above.

Example

list(ParameterGrid({"lr": [1e-3, 1e-2], "width": [8, 16]})) [{'lr': 0.001, 'width': 8}, {'lr': 0.001, 'width': 16}, {'lr': 0.01, 'width': 8}, {'lr': 0.01, 'width': 16}]

Source code in packages/pipekit-train/src/pipekit_train/sweep.py

class ParameterGrid(Operator):
    """Cartesian product over named parameter axes.

    Args:
        grid: Mapping ``axis_name → list of values``. Every axis
            must have at least one value. An **empty mapping**
            (``{}``) is also allowed and yields **zero** combinations
            (``len(grid) == 0``, ``list(grid) == []``) — this
            departs from the mathematical Cartesian-product
            convention where ``∏ over zero axes == 1``, in favour
            of the pragmatic "no axes ⇒ no trials to run" reading
            that downstream `HyperSweep` consumers expect.

    Iteration yields one ``{axis_name: value, ...}`` dict per cell.
    ``len(grid)`` is the product of axis lengths for non-empty
    grids; 0 for the empty grid per the convention above.

    Example:
        >>> list(ParameterGrid({"lr": [1e-3, 1e-2], "width": [8, 16]}))
        [{'lr': 0.001, 'width': 8}, {'lr': 0.001, 'width': 16},
         {'lr': 0.01, 'width': 8}, {'lr': 0.01, 'width': 16}]
    """

    def __init__(self, grid: dict[str, list[Any]]) -> None:
        if not isinstance(grid, dict):
            raise TypeError(
                f"ParameterGrid.grid must be a dict; got {type(grid).__name__}."
            )
        for name, values in grid.items():
            if not isinstance(values, (list, tuple)):
                raise TypeError(
                    f"ParameterGrid axis {name!r} must be a list/tuple of "
                    f"values; got {type(values).__name__}."
                )
            if len(values) == 0:
                raise ValueError(
                    f"ParameterGrid axis {name!r} is empty — every axis "
                    "must have at least one value."
                )
        self.grid = dict(grid)

    def __iter__(self) -> Iterator[dict[str, Any]]:
        if not self.grid:
            return
        axes = list(self.grid.keys())
        value_lists = [self.grid[a] for a in axes]
        for combo in itertools.product(*value_lists):
            yield dict(zip(axes, combo, strict=True))

    def __len__(self) -> int:
        if not self.grid:
            return 0
        return math.prod(len(v) for v in self.grid.values())

    def _apply(self, *args: Any, **kwargs: Any) -> Any:
        raise TypeError(
            "ParameterGrid is iterated, not called. Use `list(grid)` "
            "or pass it to `HyperSweep(search_space=...)`."
        )

SweepCarryState

Bases: CarryState

Carry-state for HyperSweep.

Threaded across trials inside _apply so the sweep composes into Sequential (e.g. coarse-then-fine sweep stages share best-so-far progress).

Attributes:

Name	Type	Description
completed		Cumulative number of trials run across all _apply calls that have consumed this state.
best_metric		Best metric value observed (None until the first trial completes).
best_trial_index		Global index (in the same coordinate system as completed) of the trial that produced best_metric. When _apply is called with an incoming SweepCarryState, trials in the new stage are numbered state.completed, state.completed + 1, … so best_trial_index always identifies the best trial unambiguously across coarse → fine sweep stages.

Source code in packages/pipekit-train/src/pipekit_train/sweep.py

class SweepCarryState(CarryState):
    """Carry-state for `HyperSweep`.

    Threaded across trials inside ``_apply`` so the sweep composes
    into `Sequential` (e.g. coarse-then-fine sweep stages share
    best-so-far progress).

    Attributes:
        completed: Cumulative number of trials run across all
            `_apply` calls that have consumed this state.
        best_metric: Best metric value observed (None until the
            first trial completes).
        best_trial_index: **Global** index (in the same coordinate
            system as ``completed``) of the trial that produced
            ``best_metric``. When `_apply` is called with an
            incoming `SweepCarryState`, trials in the new stage
            are numbered ``state.completed, state.completed + 1, …``
            so ``best_trial_index`` always identifies the best trial
            unambiguously across coarse → fine sweep stages.
    """

    def __init__(
        self,
        completed: int = 0,
        best_metric: float | None = None,
        best_trial_index: int | None = None,
    ) -> None:
        self.completed = int(completed)
        self.best_metric = best_metric if best_metric is None else float(best_metric)
        self.best_trial_index = best_trial_index

SweepResult dataclass

One trial's result inside a sweep.

Attributes:

Name	Type	Description
params	dict[str, Any]	The parameter dict that produced this trial.
trained_model_op	Operator	The trained pipekit.Operator returned by the trial's TrainingLoop.run.
artifact	Any	The trial's TrainingArtifact (or _LocalArtifact fallback when pipekit-experiment is absent).
final_metric	float	The value of select_metric pulled from artifact.backend_info["final_metrics"] — the number this trial is ranked by.
sweep_id	str	Stable identifier shared by every trial in this sweep (defaults to a uuid4 hex string).
trial_index	int	Global 0-based index across the cumulative trial stream represented by the carry-state. For a fresh sweep it starts at 0; when a HyperSweep is composed in Sequential and resumes a SweepCarryState, this index continues from state.completed at entry — so trial_index always identifies the trial unambiguously regardless of staging.

Source code in packages/pipekit-train/src/pipekit_train/sweep.py

@dataclass
class SweepResult:
    """One trial's result inside a sweep.

    Attributes:
        params: The parameter dict that produced this trial.
        trained_model_op: The trained `pipekit.Operator` returned by
            the trial's `TrainingLoop.run`.
        artifact: The trial's `TrainingArtifact` (or `_LocalArtifact`
            fallback when `pipekit-experiment` is absent).
        final_metric: The value of ``select_metric`` pulled from
            ``artifact.backend_info["final_metrics"]`` — the number
            this trial is ranked by.
        sweep_id: Stable identifier shared by every trial in this
            sweep (defaults to a uuid4 hex string).
        trial_index: **Global** 0-based index across the cumulative
            trial stream represented by the carry-state. For a fresh
            sweep it starts at 0; when a `HyperSweep` is composed in
            `Sequential` and resumes a `SweepCarryState`, this index
            continues from `state.completed` at entry — so
            `trial_index` always identifies the trial unambiguously
            regardless of staging.
    """

    params: dict[str, Any]
    trained_model_op: Operator
    artifact: Any
    final_metric: float
    sweep_id: str
    trial_index: int

Backend adapters

The Equinox+Optax+Orbax adapter is the v0.1 reference. Lightning and Keras adapter modules ship as scaffolds that raise NotImplementedError on run() until v0.2 / v0.3 land — see adapters.md for the planned designs.

Equinox (v0.1 reference)

pipekit_train.adapters.equinox

Equinox + Optax + Orbax adapter — the v0.1 reference.

Concrete JAX-native stack as described in docs/design/api/adapters.md (Equinox section). The adapter synthesises a TrainTask from a pipekit_train.Loss if the user doesn't supply one, builds the optimiser from optimizer_config, runs an eqx.filter_jit'd train step, and threads the user's callbacks through the loop.

Module surface:

• ShardingSpec — optional multi-device / multi-host sharding config (mesh, model_pspec, data_pspec); threaded through TrainState, the data iterator, and the Orbax restore path. See issue #15, ADR D12.
• TrainState(eqx.Module) — model + opt_state + step.
• TrainTask (Protocol) — user-defined training target, ported from the eqx_trainer design.
• train_step — pure function, @eqx.filter_jit wrapped.
• save_state / restore_state — Orbax bridge via eqx.partition.
• EquinoxModelOp(Operator) — wraps a trained eqx.Module as a pipekit.Operator. Deliberate in-package twin of pipekit_jax.JaxModelOp (kept separate so pipekit-train depends only on pipekit core); unlike JaxModelOp it has no weight-blob round-trip — wrap the trained module in a JaxModelOp when you need serialize_weights / from_registry.
• run(loop) — full end-to-end training entry point.

See ADRs D8, D9, D10 in docs/design/decisions.md.

EquinoxModelOp

Bases: Operator

Wrap an eqx.Module as a pipekit.Operator.

Public surface used in notebooks and downstream packages. Same constructor signature and _apply shape as pipekit_jax.JaxModelOp — swap the import when you need the registry weight round-trip (serialize_weights / with_weights / from_registry), which this in-package twin deliberately omits so pipekit-train depends only on pipekit core.

The wrapper is marked forbid_in_yaml = True because an eqx.Module is a JAX PyTree of arrays — the round-trip path is the model registry's weights blob plus the module's structural config, not JSON of the operator itself.

Parameters:

Name	Type	Description	Default
module	Module	An eqx.Module.	required

Source code in packages/pipekit-train/src/pipekit_train/adapters/equinox.py

class EquinoxModelOp(Operator):
    """Wrap an `eqx.Module` as a `pipekit.Operator`.

    Public surface used in notebooks and downstream packages. Same
    constructor signature and ``_apply`` shape as
    ``pipekit_jax.JaxModelOp`` — swap the import when you need the
    registry weight round-trip (``serialize_weights`` /
    ``with_weights`` / ``from_registry``), which this in-package twin
    deliberately omits so pipekit-train depends only on pipekit core.

    The wrapper is marked ``forbid_in_yaml = True`` because an
    ``eqx.Module`` is a JAX PyTree of arrays — the round-trip path is
    the model registry's ``weights`` blob plus the module's
    structural config, not JSON of the operator itself.

    Args:
        module: An ``eqx.Module``.
    """

    forbid_in_yaml: ClassVar[bool] = True
    __config_mixin_auto__ = False

    def __init__(self, module: eqx.Module) -> None:
        if not isinstance(module, eqx.Module):
            raise TypeError(
                "EquinoxModelOp.module must be an eqx.Module; got "
                f"{type(module).__name__}."
            )
        self.module = module

    def _apply(self, x: Any) -> Any:
        # eqx.Module is callable (instance method __call__), but ty
        # doesn't know that statically.
        return cast(Any, self.module)(x)

    def get_config(self) -> dict[str, Any]:
        # __qualname__ matches pipekit_jax.JaxModelOp so the two twins emit
        # identical configs for the same module.
        return {"module_class": type(self.module).__qualname__}

ShardingSpec dataclass

Sharding configuration for the Equinox adapter.

Threaded through TrainState (model + optimiser leaves placed per model_pspec), the data iterator (each batch placed per data_pspec, with multi-host process sharding), and the Orbax restore path (sharding-aware). See ADR D12.

The common case is data parallelism: model_pspec = PartitionSpec() (the model is replicated on every device) and data_pspec = PartitionSpec("data") (each batch is split along its leading axis across the mesh's "data" dimension). A uniform model-parallel model_pspec is also honoured when every sharded array leaf has a compatible rank.

Attributes:

Name	Type	Description
mesh	Mesh	The device mesh.
model_pspec	PartitionSpec	PartitionSpec applied to every model / optimiser array leaf. PartitionSpec() replicates.
data_pspec	PartitionSpec	PartitionSpec applied to each batch (the leading axis is the batch axis).

Source code in packages/pipekit-train/src/pipekit_train/adapters/equinox.py

@dataclass(frozen=True)
class ShardingSpec:
    """Sharding configuration for the Equinox adapter.

    Threaded through `TrainState` (model + optimiser leaves placed per
    ``model_pspec``), the data iterator (each batch placed per
    ``data_pspec``, with multi-host process sharding), and the Orbax
    restore path (sharding-aware). See ADR D12.

    The common case is **data parallelism**: ``model_pspec =
    PartitionSpec()`` (the model is replicated on every device) and
    ``data_pspec = PartitionSpec("data")`` (each batch is split along its
    leading axis across the mesh's ``"data"`` dimension). A uniform
    model-parallel ``model_pspec`` is also honoured when every sharded
    array leaf has a compatible rank.

    Attributes:
        mesh: The device mesh.
        model_pspec: PartitionSpec applied to every model / optimiser
            array leaf. ``PartitionSpec()`` replicates.
        data_pspec: PartitionSpec applied to each batch (the leading axis
            is the batch axis).
    """

    mesh: jax.sharding.Mesh
    model_pspec: jax.sharding.PartitionSpec
    data_pspec: jax.sharding.PartitionSpec

    @property
    def model_sharding(self) -> jax.sharding.NamedSharding:
        """`NamedSharding` for model / optimiser leaves."""
        return jax.sharding.NamedSharding(self.mesh, self.model_pspec)

    @property
    def data_sharding(self) -> jax.sharding.NamedSharding:
        """`NamedSharding` for data batches."""
        return jax.sharding.NamedSharding(self.mesh, self.data_pspec)

    @property
    def replicated(self) -> jax.sharding.NamedSharding:
        """`NamedSharding` that replicates a value on every device."""
        return jax.sharding.NamedSharding(self.mesh, jax.sharding.PartitionSpec())

data_sharding property

NamedSharding for data batches.

model_sharding property

NamedSharding for model / optimiser leaves.

replicated property

NamedSharding that replicates a value on every device.

TrainState

Bases: Module

Equinox-side training state.

A PyTree-clean container holding the bare eqx.Module (without the EquinoxModelOp wrapper), the Optax optimizer state, and the global step counter.

Source code in packages/pipekit-train/src/pipekit_train/adapters/equinox.py

class TrainState(eqx.Module):
    """Equinox-side training state.

    A PyTree-clean container holding the bare ``eqx.Module``
    (without the ``EquinoxModelOp`` wrapper), the Optax optimizer
    state, and the global step counter.
    """

    model: eqx.Module
    opt_state: optax.OptState
    step: jax.Array

    @staticmethod
    def create(
        model: eqx.Module,
        optimizer: optax.GradientTransformation,
        sharding: ShardingSpec | None = None,
    ) -> TrainState:
        """Build the initial training state, optionally placed on a mesh.

        When ``sharding`` is given, the model and optimiser array leaves are
        placed on ``sharding.model_sharding`` and the step counter is
        replicated, so the very first ``train_step`` already runs sharded.
        """
        opt_state = optimizer.init(eqx.filter(model, eqx.is_array))
        state = TrainState(
            model=model,
            opt_state=opt_state,
            step=jnp.asarray(0, dtype=jnp.int32),
        )
        if sharding is not None:
            state = shard_train_state(state, sharding)
        return state

create(model, optimizer, sharding=None) staticmethod

Build the initial training state, optionally placed on a mesh.

When sharding is given, the model and optimiser array leaves are placed on sharding.model_sharding and the step counter is replicated, so the very first train_step already runs sharded.

Source code in packages/pipekit-train/src/pipekit_train/adapters/equinox.py

@staticmethod
def create(
    model: eqx.Module,
    optimizer: optax.GradientTransformation,
    sharding: ShardingSpec | None = None,
) -> TrainState:
    """Build the initial training state, optionally placed on a mesh.

    When ``sharding`` is given, the model and optimiser array leaves are
    placed on ``sharding.model_sharding`` and the step counter is
    replicated, so the very first ``train_step`` already runs sharded.
    """
    opt_state = optimizer.init(eqx.filter(model, eqx.is_array))
    state = TrainState(
        model=model,
        opt_state=opt_state,
        step=jnp.asarray(0, dtype=jnp.int32),
    )
    if sharding is not None:
        state = shard_train_state(state, sharding)
    return state

TrainTask

Bases: Protocol

User-defined Equinox training target.

Provide loss_fn (required); optionally provide eval_fn. The adapter synthesises a TrainTask from loop.loss when no explicit task is supplied.

Signatures

loss_fn(model, batch, key) -> (scalar_loss, aux_metrics_dict) eval_fn(model, batch, key) -> metrics_dict (optional)

Source code in packages/pipekit-train/src/pipekit_train/adapters/equinox.py

@runtime_checkable
class TrainTask(Protocol):
    """User-defined Equinox training target.

    Provide ``loss_fn`` (required); optionally provide ``eval_fn``.
    The adapter synthesises a `TrainTask` from ``loop.loss`` when no
    explicit task is supplied.

    Signatures:
        loss_fn(model, batch, key) -> (scalar_loss, aux_metrics_dict)
        eval_fn(model, batch, key) -> metrics_dict  (optional)
    """

    def loss_fn(
        self,
        model: eqx.Module,
        batch: Any,
        key: jax.Array,
    ) -> tuple[jax.Array, dict[str, jax.Array]]: ...

restore_state(mngr, template, step)

Restore array leaves into the static skeleton of template.

Returns:

Type	Description
TrainState	(restored_state, data_iter_state). data_iter_state is
Any | None	None when no side-car was written at the checkpoint step
tuple[TrainState, Any | None]	(e.g. a checkpoint written by an earlier pipekit-train
tuple[TrainState, Any | None]	without the iterator-state extension). When non-None the
tuple[TrainState, Any | None]	caller passes it to grain_iterator.set_state(state) to
tuple[TrainState, Any | None]	resume the data stream byte-identically. See #17.

Source code in packages/pipekit-train/src/pipekit_train/adapters/equinox.py

def restore_state(
    mngr: ocp.CheckpointManager,
    template: TrainState,
    step: int,
) -> tuple[TrainState, Any | None]:
    """Restore array leaves into the static skeleton of ``template``.

    Returns:
        ``(restored_state, data_iter_state)``. ``data_iter_state`` is
        ``None`` when no side-car was written at the checkpoint step
        (e.g. a checkpoint written by an earlier pipekit-train
        without the iterator-state extension). When non-``None`` the
        caller passes it to ``grain_iterator.set_state(state)`` to
        resume the data stream byte-identically. See #17.
    """
    import json
    import os
    from pathlib import Path

    import orbax.checkpoint as ocp  # ty: ignore[unresolved-import]

    arrays, static = eqx.partition(template, eqx.is_array)
    abstract = jax.tree.map(ocp.utils.to_shape_dtype_struct, arrays)
    restored = mngr.restore(step, args=ocp.args.StandardRestore(abstract))

    # Read this process's own iterator-state side-car (see save_state). Fall
    # back to the legacy shared name for checkpoints written before the
    # per-process split.
    step_dir = Path(os.fspath(mngr.directory)) / str(step)
    side_car = step_dir / f"data_iter_{jax.process_index()}.json"
    if not side_car.exists():
        side_car = step_dir / "data_iter.json"
    data_iter_state: Any | None = None
    if side_car.exists():
        data_iter_state = _unjsonify_grain_state(json.loads(side_car.read_text()))

    return (
        cast(TrainState, eqx.combine(restored, static)),
        data_iter_state,
    )

run(loop)

Train loop end-to-end. The entry point for the Equinox backend.

Returns:

Type	Description
Operator	A pair (trained_model_op, backend_info). trained_model_op
dict[str, Any]	is a EquinoxModelOp wrapping the trained eqx.Module;
tuple[Operator, dict[str, Any]]	backend_info is a dict with the keys backend,
tuple[Operator, dict[str, Any]]	jax_version, equinox_version, optax_version,
tuple[Operator, dict[str, Any]]	devices, total_seconds, final_step, and
tuple[Operator, dict[str, Any]]	final_metrics for inclusion in the TrainingArtifact.

Source code in packages/pipekit-train/src/pipekit_train/adapters/equinox.py

def run(loop: TrainingLoop) -> tuple[Operator, dict[str, Any]]:
    """Train ``loop`` end-to-end. The entry point for the Equinox backend.

    Returns:
        A pair ``(trained_model_op, backend_info)``. ``trained_model_op``
        is a `EquinoxModelOp` wrapping the trained ``eqx.Module``;
        ``backend_info`` is a dict with the keys ``backend``,
        ``jax_version``, ``equinox_version``, ``optax_version``,
        ``devices``, ``total_seconds``, ``final_step``, and
        ``final_metrics`` for inclusion in the TrainingArtifact.
    """
    # Lazy imports — keep this module loadable without the full
    # [equinox] extra installed (see module docstring).
    try:
        import optax  # ty: ignore[unresolved-import]
        import orbax.checkpoint as ocp  # ty: ignore[unresolved-import]
    except ImportError as exc:
        raise ImportError(
            "backend='equinox' needs the full [equinox] extra (optax + "
            "orbax-checkpoint + grain). Install with "
            "`pip install pipekit-train[equinox]`."
        ) from exc

    t0 = time.time()
    # --- Unwrap the model ----------------------------------------------
    model_op = loop.model_op
    if isinstance(model_op, EquinoxModelOp):
        eqx_module = model_op.module
    elif isinstance(model_op, eqx.Module):
        eqx_module = model_op
    elif hasattr(model_op, "module") and isinstance(model_op.module, eqx.Module):
        # pipekit_jax.JaxModelOp (or anything module-shaped) lands here.
        eqx_module = model_op.module
    else:
        raise TypeError(
            "Equinox adapter expects loop.model_op to be a pipekit_train."
            "adapters.equinox.EquinoxModelOp (or any operator exposing "
            ".module: eqx.Module). Got "
            f"{type(model_op).__name__}."
        )

    # --- Task / loss validation ----------------------------------------
    task: Any = loop.task
    if task is None:
        if loop.loss is None:
            raise ValueError(
                "Equinox adapter requires loop.loss or loop.task to be "
                "supplied. Pass a Loss (carrier-agnostic) or your own "
                "TrainTask (loss_fn-shaped) when constructing TrainingLoop."
            )
        task = _SynthesisedTask(loop.loss)
    elif not isinstance(task, TrainTask):
        raise TypeError(
            "Equinox adapter: loop.task must satisfy the TrainTask Protocol "
            f"(missing `loss_fn`); got {type(task).__name__}."
        )

    # --- Sharding (optional; issue #15) --------------------------------
    sharding: ShardingSpec | None = getattr(loop, "sharding", None)
    if sharding is not None and not isinstance(sharding, ShardingSpec):
        raise TypeError(
            "Equinox adapter: loop.sharding must be a pipekit_train.adapters."
            f"equinox.ShardingSpec or None; got {type(sharding).__name__}."
        )

    # --- Optimiser, RNG -------------------------------------------------
    optimizer = _build_optimizer(loop.optimizer_config)

    state = TrainState.create(eqx_module, optimizer, sharding=sharding)
    rng = jax.random.key(loop.seed)

    # --- Checkpoint manager (if checkpoint_dir set) --------------------
    mngr: ocp.CheckpointManager | None = None
    if loop.checkpoint_dir:
        # Look for a Checkpoint callback to learn cadence + retention.
        every_n = loop.eval_every_n_steps
        keep_last = 3
        for cb in loop.callbacks:
            if hasattr(cb, "every_n_steps") and hasattr(cb, "keep_last"):
                every_n = int(cb.every_n_steps)
                keep_last = int(cb.keep_last)
                break
        mngr = ocp.CheckpointManager(
            directory=os.path.abspath(loop.checkpoint_dir),
            options=ocp.CheckpointManagerOptions(
                max_to_keep=keep_last,
                save_interval_steps=every_n,
            ),
        )
    # Resume from the latest checkpoint if one exists. Restores both
    # TrainState (weights + opt_state) AND the iterator position.
    # The data-iter state is saved aside for after we build the
    # iterator below.
    data_iter_state_to_restore: Any = None
    if mngr is not None:
        latest = mngr.latest_step()
        if latest is not None:
            state, data_iter_state_to_restore = restore_state(mngr, state, latest)

    # --- Data iterator -------------------------------------------------
    # Built only if we're actually going to consume batches. Skipping
    # this for finished resumes (state.step >= max_steps) means
    # users can reload a checkpointed model on a dataset smaller than
    # batch_size without tripping `_BatchSource`'s indexable-min-size
    # guard. See #17 / codex review on PR #24.
    train_iter: _BatchSource | None = None
    if int(state.step) < loop.max_steps:
        train_iter = _BatchSource(
            loop.dataset, loop.batch_size, loop.seed, sharding=sharding
        )
        if data_iter_state_to_restore is not None:
            train_iter.set_state(data_iter_state_to_restore)

    # --- Initial-state callback dispatch -------------------------------
    initial_state = _carry_state(state, loop)
    _dispatch(loop.callbacks, "on_train_begin", loop, initial_state)

    # --- Main loop -----------------------------------------------------
    # try/finally: a raising callback / loss / NaN guard must not leak
    # the metric writer's file handle (dropping buffered lines) or leave
    # an async checkpoint half-written.
    last_metrics: dict[str, float] = {}
    try:
        # Inside the loop body, `train_iter` is guaranteed non-None because
        # the loop condition matches the guard that built it.
        while int(state.step) < loop.max_steps:
            assert train_iter is not None  # see iterator-build guard above
            rng, step_key = jax.random.split(rng)
            batch = train_iter.next_batch()
            state, metrics = train_step(state, batch, step_key, task, optimizer)
            last_metrics = {k: float(np.asarray(v)) for k, v in metrics.items()}

            step = int(state.step)

            # Per-step log to the writer (if any).
            if loop.metric_writer is not None and step % loop.log_every_n_steps == 0:
                loop.metric_writer.write(step, last_metrics)

            # Callbacks on_step_end.
            carry = _carry_state(state, loop, last_metrics)
            _dispatch(loop.callbacks, "on_step_end", loop, carry, last_metrics)

            # Periodic epoch boundary (cosmetic; controls on_epoch_end firing).
            if loop.steps_per_epoch and step > 0 and step % loop.steps_per_epoch == 0:
                _dispatch(loop.callbacks, "on_epoch_end", loop, carry, last_metrics)

            # Periodic evaluation.
            if (
                loop.val_dataset is not None
                and step > 0
                and step % loop.eval_every_n_steps == 0
            ):
                rng, eval_key = jax.random.split(rng)
                eval_metrics = _evaluate(
                    state.model, loop.val_dataset, loop.batch_size, task, eval_key
                )
                _dispatch(loop.callbacks, "on_eval_end", loop, carry, eval_metrics)

            # Checkpoint — gated on the manager's configured cadence
            # (CheckpointManager.save_interval_steps, derived from the
            # Checkpoint callback's every_n_steps). Calling save() every
            # step would defeat the configured cadence and add I/O on
            # the fast path.
            if mngr is not None and step > 0 and mngr.should_save(step):
                save_state(mngr, state, step, data_iter_state=train_iter.get_state())

            # Early-stopping check — break before next step.
            if _any_should_stop(loop.callbacks):
                break

        # --- Final eval (if val_dataset and we didn't just eval) ------
        if (
            loop.val_dataset is not None
            and int(state.step) % loop.eval_every_n_steps != 0
        ):
            rng, eval_key = jax.random.split(rng)
            final_eval = _evaluate(
                state.model, loop.val_dataset, loop.batch_size, task, eval_key
            )
            carry = _carry_state(state, loop, last_metrics)
            _dispatch(loop.callbacks, "on_eval_end", loop, carry, final_eval)
    finally:
        # --- Wait for any async checkpoint + close the writer ---------
        if mngr is not None:
            mngr.wait_until_finished()
        if loop.metric_writer is not None:
            loop.metric_writer.close()

    # --- Wrap the trained module + dispatch on_train_end -------------
    trained_model_op = EquinoxModelOp(state.model)
    final_carry = _carry_state(
        state, loop, last_metrics, model_override=trained_model_op
    )
    _dispatch(loop.callbacks, "on_train_end", loop, final_carry)

    backend_info = {
        "backend": "equinox",
        "jax_version": jax.__version__,
        "equinox_version": eqx.__version__,
        "optax_version": optax.__version__,
        "devices": [str(d) for d in jax.devices()],
        "total_seconds": time.time() - t0,
        "final_step": int(state.step),
        "final_metrics": last_metrics,
        "sharding": _sharding_info(sharding),
    }
    return trained_model_op, backend_info

save_state(mngr, state, step, *, data_iter_state=None)

Persist the array leaves of state at step.

Static (non-array) leaves are reconstructed on restore from the in-memory template; see restore_state.

Parameters:

Name	Type	Description	Default
mngr	CheckpointManager	Orbax checkpoint manager.	required
state	TrainState	Current TrainState.	required
step	int	Step number.	required
data_iter_state	Any	Optional state dict from grain_iterator.get_state(). When provided, it's persisted as a JSON side-car next to the Orbax checkpoint (<directory>/<step>/data_iter.json) so resume-from- checkpoint restores byte-identical data-iter position alongside the weights. See #17.	None

Source code in packages/pipekit-train/src/pipekit_train/adapters/equinox.py

def save_state(
    mngr: ocp.CheckpointManager,
    state: TrainState,
    step: int,
    *,
    data_iter_state: Any = None,
) -> None:
    """Persist the array leaves of ``state`` at ``step``.

    Static (non-array) leaves are reconstructed on restore from the
    in-memory template; see ``restore_state``.

    Args:
        mngr: Orbax checkpoint manager.
        state: Current ``TrainState``.
        step: Step number.
        data_iter_state: Optional state dict from
            ``grain_iterator.get_state()``. When provided, it's
            persisted as a JSON side-car next to the Orbax checkpoint
            (``<directory>/<step>/data_iter.json``) so resume-from-
            checkpoint restores byte-identical data-iter position
            alongside the weights. See #17.
    """
    import json
    import os
    from pathlib import Path

    import orbax.checkpoint as ocp  # ty: ignore[unresolved-import]

    arrays, _ = eqx.partition(state, eqx.is_array)
    mngr.save(step, args=ocp.args.StandardSave(arrays))

    if data_iter_state is not None:
        # Orbax writes the checkpoint asynchronously; the per-step dir
        # exists after wait_until_finished. For correctness, write the
        # side-car eagerly into the same path. Orbax exposes
        # `mngr.directory / str(step)` as the per-step layout.
        mngr.wait_until_finished()
        step_dir = Path(os.fspath(mngr.directory)) / str(step)
        step_dir.mkdir(parents=True, exist_ok=True)
        # Per-process side-car: under multi-host sharding each process owns a
        # distinct data shard and so a distinct Grain iterator state. A single
        # shared file would let the last writer win and every process would
        # then resume from the wrong position. Key it by process index.
        (step_dir / f"data_iter_{jax.process_index()}.json").write_text(
            json.dumps(_jsonify_grain_state(data_iter_state), sort_keys=True)
        )

shard_train_state(state, spec)

Return state with its model + optimiser leaves placed on spec.

The model and optimiser arrays go to spec.model_sharding; the scalar step is replicated across the mesh.

Source code in packages/pipekit-train/src/pipekit_train/adapters/equinox.py

def shard_train_state(state: TrainState, spec: ShardingSpec) -> TrainState:
    """Return ``state`` with its model + optimiser leaves placed on ``spec``.

    The model and optimiser arrays go to ``spec.model_sharding``; the scalar
    step is replicated across the mesh.
    """
    return TrainState(
        model=_shard_pytree(state.model, spec.model_sharding, spec.replicated),
        opt_state=_shard_pytree(state.opt_state, spec.model_sharding, spec.replicated),
        step=jax.device_put(state.step, spec.replicated),
    )

train_step(state, batch, key, task, optimizer)

One Equinox optimisation step. Pure function.

Differentiates only array leaves (eqx.filter_grad); passes task and optimizer as static arguments so structural changes recompile.

Source code in packages/pipekit-train/src/pipekit_train/adapters/equinox.py

@eqx.filter_jit
def train_step(
    state: TrainState,
    batch: Any,
    key: jax.Array,
    task: TrainTask,
    optimizer: optax.GradientTransformation,
) -> tuple[TrainState, dict[str, jax.Array]]:
    """One Equinox optimisation step. Pure function.

    Differentiates only array leaves (`eqx.filter_grad`); passes
    ``task`` and ``optimizer`` as static arguments so structural
    changes recompile.
    """

    @eqx.filter_grad(has_aux=True)
    def grad_fn(
        model: eqx.Module,
        batch: Any,
        key: jax.Array,
    ) -> tuple[jax.Array, dict[str, jax.Array]]:
        return task.loss_fn(model, batch, key)

    grads, metrics = grad_fn(state.model, batch, key)
    params = eqx.filter(state.model, eqx.is_array)
    updates, new_opt_state = optimizer.update(grads, state.opt_state, params)
    new_model = eqx.apply_updates(state.model, updates)
    return (
        TrainState(model=new_model, opt_state=new_opt_state, step=state.step + 1),
        metrics,
    )