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; a Keras adapter module is scaffolded for v0.3). Trained models are first-class pipekit_ops.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'
in_memory	bool	Load the whole cache into NumPy arrays on first read and serve items from them. A zarr read costs about 1 ms per array, so turn this on when the cache fits in memory and the dataset is read item by item, as in training.	False

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.
• The store is opened once per instance and process, and reopened after invalidate or a change of content hash. Delete the cache through invalidate: an instance holding the arrays does not see a directory removed behind its back.

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.
        in_memory: Load the whole cache into NumPy arrays on first read
            and serve items from them. A zarr read costs about 1 ms per
            array, so turn this on when the cache fits in memory and the
            dataset is read item by item, as in training.

    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.
        - The store is opened once per instance and process, and reopened
          after `invalidate` or a change of content hash. Delete the cache
          through `invalidate`: an instance holding the arrays does not see
          a directory removed behind its back.
    """

    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",
        in_memory: bool = False,
    ) -> 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
        self.in_memory = in_memory
        self._arrays: tuple[str, Any, Any] | None = None

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

    def __getstate__(self) -> dict[str, Any]:
        # Each worker process opens the store for itself.
        return {**self.__dict__, "_arrays": None}

    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 _cached_arrays(self) -> tuple[Any, Any]:
        """The ``x`` and ``y`` arrays, opening the store at most once per hash."""
        path = self._cache_path()
        if self._arrays is None or self._arrays[0] != path:
            if not self._materialised():
                self._materialise()
            x_arr, y_arr = self._open_cache()
            if self.in_memory:
                x_arr, y_arr = x_arr[:], y_arr[:]
            self._arrays = (path, x_arr, y_arr)
        return self._arrays[1], self._arrays[2]

    def __getitem__(self, idx: int) -> tuple[Any, Any]:
        x_arr, y_arr = self._cached_arrays()
        return x_arr[idx], y_arr[idx]

    def __iter__(self) -> Iterator[tuple[Any, Any]]:
        x_arr, y_arr = self._cached_arrays()
        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.
        """
        self._arrays = None
        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,
            "in_memory": self.in_memory,
        }

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.
    """
    self._arrays = None
    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	Operator returning a parameter sample. It is called with seed= if its signature accepts that, and with no arguments otherwise; the choice is made once, at construction. Observation noise drawn inside the prior is seeded by this call.	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. A stateful cycle is called as cycle(params, state), and the dataset keeps the carrier it returns; any other operator as cycle(params).	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'
cycle_state	CarryState | None	Initial carry state for a stateful cycle, such as a pipekit_cycle.DAState. Defaults to a new CarryState. Each item gets its own deep copy, so samples don't depend on the order they're read in. Part of the content hash when set.	None
simulator_version	str | None	Label for the simulator's behaviour, such as "qg-v2". Change it whenever the simulator changes in a way its config can't show, because a Python function can't be hashed reliably.	None
invert	bool	False yields (params, output), for training an emulator. True yields (output, params), for training an amortized posterior that maps an observation to parameters. Part of the content hash when set, because the two orders must not share a cache.	False

The content hash covers the class and get_config() of the forward model, cycle and observation operator, the forward model's state_signature and dt, the prior config, n_samples, seed, split, cycle_state and simulator_version. Code changes that none of these show aren't detected. content_hash warns once per dataset when the forward model is a wrapped function or has no get_config, and no simulator_version is set.

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: `Operator` returning a parameter sample. It is called with
            ``seed=`` if its signature accepts that, and with no arguments
            otherwise; the choice is made once, at construction. Observation
            noise drawn inside the prior is seeded by this call.
        n_samples: Total samples per epoch.
        cycle: Optional ``pipekit_cycle.Cycle`` for trajectory rollouts.
            When ``None``, the dataset emits single-step pairs. A stateful
            cycle is called as ``cycle(params, state)``, and the dataset
            keeps the carrier it returns; any other operator as
            ``cycle(params)``.
        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).
        cycle_state: Initial carry state for a stateful ``cycle``, such as a
            ``pipekit_cycle.DAState``. Defaults to a new ``CarryState``. Each
            item gets its own deep copy, so samples don't depend on the order
            they're read in. Part of the content hash when set.
        simulator_version: Label for the simulator's behaviour, such as
            ``"qg-v2"``. Change it whenever the simulator changes in a way its
            config can't show, because a Python function can't be hashed
            reliably.
        invert: ``False`` yields ``(params, output)``, for training an emulator.
            ``True`` yields ``(output, params)``, for training an amortized
            posterior that maps an observation to parameters. Part of the
            content hash when set, because the two orders must not share a cache.

    The content hash covers the class and ``get_config()`` of the forward
    model, cycle and observation operator, the forward model's
    ``state_signature`` and ``dt``, the prior config, ``n_samples``, ``seed``,
    ``split``, ``cycle_state`` and ``simulator_version``. Code changes that
    none of these show aren't detected. `content_hash` warns once per dataset
    when the forward model is a wrapped function or has no ``get_config``, and
    no ``simulator_version`` is set.
    """

    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",
        *,
        cycle_state: CarryState | None = None,
        simulator_version: str | None = None,
        invert: bool = False,
    ) -> 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_ops.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._prior_accepts_seed = accepts_keyword(prior, "seed")
        self.n_samples = int(n_samples)
        self.cycle = cycle
        self.cycle_state = cycle_state
        self.simulator_version = simulator_version
        self.invert = invert
        self._warned_unidentified = False
        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
        if self._prior_accepts_seed:
            params = self.prior(seed=sample_seed)
        else:
            params = self.prior()
        if self.cycle is not None:
            output = self._roll_out(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 (output, params) if self.invert else (params, output)

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

    def _roll_out(self, cycle: Operator, params: Any) -> Any:
        """Run ``cycle`` from ``params`` and return the final carrier."""
        if not getattr(cycle, "_is_stateful", False):
            return cycle(params)
        if self.cycle_state is None:
            state = CarryState()
        else:
            state = copy.deepcopy(self.cycle_state)
        output, _ = cycle(params, state)
        return output

    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 _warn_if_simulator_unidentified(self) -> None:
        """Warn once if the hash can't tell this simulator from another."""
        if self._warned_unidentified or self.simulator_version is not None:
            return
        wraps_function = callable(getattr(self.forward_model, "fn", None))
        has_config = callable(getattr(self.forward_model, "get_config", None))
        if has_config and not wraps_function:
            return
        self._warned_unidentified = True
        warnings.warn(
            f"SimulationDataset can't tell this {type(self.forward_model).__name__} "
            "from another of the same class, so a cache may serve stale outputs "
            "after the simulator changes. Pass simulator_version= to identify it.",
            UserWarning,
            stacklevel=3,
        )

    def content_hash(self) -> str:
        self._warn_if_simulator_unidentified()
        fm_sig = getattr(self.forward_model, "state_signature", None)
        fm_dt = getattr(self.forward_model, "dt", None)
        recipe: dict[str, Any] = {
            "forward_model": {
                **cast(dict[str, Any], _component_identity(self.forward_model)),
                "state_signature": repr(fm_sig),
                "dt": fm_dt,
                "version": self.simulator_version,
            },
            "prior": {
                "class": type(self.prior).__qualname__,
                "config": self.prior.get_config(),
            },
            "n_samples": self.n_samples,
            "cycle": _component_identity(self.cycle),
            "obs_op": _component_identity(self.obs_op),
            "seed": self.seed,
            "split": self.split,
        }
        if self.cycle_state is not None:
            recipe["cycle_state"] = _content_key(self.cycle_state)
        # Only when set, so that the default keeps its hash.
        if self.invert:
            recipe["invert"] = True
        return hash_config("SimulationDataset", recipe)

    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
            ),
            "cycle_state": (
                stable_repr(self.cycle_state) if self.cycle_state is not None else None
            ),
            "simulator_version": self.simulator_version,
            "invert": self.invert,
            "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_ops.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_ops.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 estimates it with samples drawn from predicted, which needs a PRNG key.
• 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. A component that takes a key gets its own, split from the one passed in.

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. A component that takes a ``key`` gets
    its own, split from the one passed in.

    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,
        key: Any = None,
    ) -> tuple[Any, dict[str, Any]]:
        aux: dict[str, Any] = {}
        total = None
        keys: Any = [None] * len(self.components)
        if key is not None:
            import jax

            keys = jax.random.split(key, len(self.components))
        for i, (weight, loss) in enumerate(self.components):
            result = call_loss(loss, predicted, target, keys[i])
            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.

KL(q ‖ p) = E_q[log q(z) - log p(z)], where q is predicted and p is target. The analytic path is used when predicted has .kl_divergence(target). Otherwise it is estimated by Monte Carlo: (1 / S) Σ_s [log q(z_s) - log p(z_s)], with S draws z_s from q.

Algorithm

1. Use predicted.kl_divergence(target) if it exists.
2. Otherwise draw n_samples from predicted with key.
3. Average log_prob differences over draws and any batch dims.

Gradients flow through the draws only if predicted samples by reparameterisation.

Parameters:

Name	Type	Description	Default
n_samples	int	Monte Carlo draws S for the estimate. Its error falls as 1 / √S.	32

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

Raises:

Type	Description
ValueError	If n_samples is below 1.
TypeError	If the estimate is needed and no key was given.

Examples:

>>> class Point:
...     def kl_divergence(self, other):
...         import numpy as np
...
...         return np.array([0.5, 1.5])
>>> KL()(Point(), None)[0].item()
1.0

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

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

    KL(q ‖ p) = E_q[log q(z) - log p(z)], where q is ``predicted`` and p is
    ``target``. The analytic path is used when ``predicted`` has
    ``.kl_divergence(target)``. Otherwise it is estimated by Monte Carlo:
    (1 / S) Σ_s [log q(z_s) - log p(z_s)], with S draws z_s from q.

    Algorithm:
        1. Use ``predicted.kl_divergence(target)`` if it exists.
        2. Otherwise draw ``n_samples`` from ``predicted`` with ``key``.
        3. Average ``log_prob`` differences over draws and any batch dims.

    Gradients flow through the draws only if ``predicted`` samples by
    reparameterisation.

    Args:
        n_samples: Monte Carlo draws S for the estimate. Its error falls as
            1 / √S.

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

    Raises:
        ValueError: If ``n_samples`` is below 1.
        TypeError: If the estimate is needed and no ``key`` was given.

    Examples:
        >>> class Point:
        ...     def kl_divergence(self, other):
        ...         import numpy as np
        ...
        ...         return np.array([0.5, 1.5])
        >>> KL()(Point(), None)[0].item()
        1.0
    """

    def __init__(self, n_samples: int = 32) -> None:
        if n_samples < 1:
            raise ValueError(f"KL.n_samples must be >= 1; got {n_samples}.")
        self.n_samples = n_samples

    def _apply(
        self,
        predicted: Any,
        target: Any,
        key: Any = None,
    ) -> tuple[Any, dict[str, Any]]:
        if hasattr(predicted, "kl_divergence"):
            loss = predicted.kl_divergence(target).mean()
            return loss, {"kl": loss}
        if key is None:
            raise TypeError(
                "Sample-based KL requires a PRNG key. The Equinox adapter "
                "supplies one per batch; for ad-hoc use, pass "
                "key=jax.random.key(seed)."
            )
        # z_s ~ q
        samples = _draw_samples(predicted, key, self.n_samples)
        # log q(z_s)
        log_q = predicted.log_prob(samples)
        # log p(z_s)
        log_p = target.log_prob(samples)
        # KL(q ‖ p) ≈ (1 / S) Σ_s [log q(z_s) - log p(z_s)]
        loss = (log_q - log_p).mean()
        return loss, {"kl": loss}

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. A loss that needs randomness may also take a keyword key, which the adapter supplies per batch; see call_loss.

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. A loss that needs randomness may also take a keyword
    ``key``, which the adapter supplies per batch; see `call_loss`.

    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}

call_loss(loss, predicted, target, key=None)

Call loss, passing key only if it takes one.

Parameters:

Name	Type	Description	Default
loss	Any	Any Loss-shaped callable.	required
predicted	Any	The model's prediction.	required
target	Any	The target.	required
key	Any	Optional PRNG key.	None

Returns:

Type	Description
Any	Whatever loss returns.

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

def call_loss(loss: Any, predicted: Any, target: Any, key: Any = None) -> Any:
    """Call ``loss``, passing ``key`` only if it takes one.

    Args:
        loss: Any `Loss`-shaped callable.
        predicted: The model's prediction.
        target: The target.
        key: Optional PRNG key.

    Returns:
        Whatever ``loss`` returns.
    """
    if key is not None and accepts_keyword(loss, "key"):
        return loss(predicted, target, key=key)
    return loss(predicted, target)

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.

The Equinox adapter saves into TrainingLoop.checkpoint_dir or save_dir on this cadence, and always saves the final step. (The actual save call lives in the adapter; this callback only configures it.) Only the Equinox backend writes on-disk snapshots.

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).	3
save_dir	str | PathLike[str] | None	Local directory for snapshots, used when TrainingLoop.checkpoint_dir isn't set. Setting the two to different directories raises ValueError. With neither set, nothing is written to disk, which suits a callback that only uses registry.	None
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.

    The Equinox adapter saves into ``TrainingLoop.checkpoint_dir`` or
    ``save_dir`` on this cadence, and always saves the final step. (The
    actual save call lives in the adapter; this callback only configures it.)
    Only the Equinox backend writes on-disk snapshots.

    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`).
        save_dir: Local directory for snapshots, used when
            ``TrainingLoop.checkpoint_dir`` isn't set. Setting the two to
            different directories raises ``ValueError``. With neither set,
            nothing is written to disk, which suits a callback that only uses
            ``registry``.
        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 | os.PathLike[str] | None = None,
        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 = None if save_dir is None else os.fspath(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_ops.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_ops.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). 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). 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/keras module wrapped in a pipekit_ops.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', 'keras', 'numpyro-svi', 'numpyro-mcmc', 'blackjax']	One of "equinox", "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/keras module wrapped in a
            pipekit_ops.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"``, ``"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",
            "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

        if not trained_model_hash:
            trained_model_hash = _trained_model_hash(trained_model_op)

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

        # `pipekit_ops.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.

Trials rank on validation metrics when they have them: the backend's final_eval_metrics, scored on the final state. Without a val_dataset they rank on the last training batch, with a warning.

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 the chosen split's 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
select_split	Literal['auto', 'train', 'eval']	Which metrics to rank on. "auto" (default) uses validation metrics when a trial has any, else training metrics with one warning per sweep. "eval" requires validation metrics, and "train" always uses the last training batch.	'auto'

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``.

    Trials rank on validation metrics when they have them: the backend's
    ``final_eval_metrics``, scored on the final state. Without a
    ``val_dataset`` they rank on the last training batch, with a warning.

    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 the chosen split's 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.
        select_split: Which metrics to rank on. ``"auto"`` (default) uses
            validation metrics when a trial has any, else training metrics
            with one warning per sweep. ``"eval"`` requires validation metrics,
            and ``"train"`` always uses the last training batch.

    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,
        select_split: Literal["auto", "train", "eval"] = "auto",
    ) -> 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 select_split not in {"auto", "train", "eval"}:
            raise ValueError(
                "HyperSweep.select_split must be 'auto', 'train' or 'eval'; "
                f"got {select_split!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.select_split = select_split
        self._warned_training_fallback = False
        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()
        self._warned_training_fallback = False
        # 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 {}
        key = self._metrics_key(backend_info)
        metrics = backend_info.get(key, {}) or {}
        if self.select_metric not in metrics:
            raise ValueError(
                f"HyperSweep: trial {trial_index} (params={params!r}) "
                f"artifact has no metric {self.select_metric!r} in "
                f"backend_info[{key!r}]. Available keys: "
                f"{sorted(metrics)}."
            )
        return float(metrics[self.select_metric])

    def _metrics_key(self, backend_info: dict[str, Any]) -> str:
        """Return the ``backend_info`` key holding the metrics to rank on."""
        if self.select_split == "train":
            return "final_metrics"
        if self.select_split == "eval" or backend_info.get("final_eval_metrics"):
            return "final_eval_metrics"
        if not self._warned_training_fallback:
            self._warned_training_fallback = True
            warnings.warn(
                "HyperSweep is ranking on the last training batch, because a trial "
                "has no validation metrics. Give the loop a val_dataset, or pass "
                "select_split='train' to choose this explicitly.",
                UserWarning,
                stacklevel=2,
            )
        return "final_metrics"

    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,
            "select_split": self.select_split,
            "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_ops.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 this trial is ranked by, from the split that HyperSweep.select_split chose.
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_ops.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`` this trial is ranked
            by, from the split that ``HyperSweep.select_split`` chose.
        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. The Keras adapter module ships as a scaffold that raises NotImplementedError on run() until v0.3 lands — 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_ops.Operator. Deliberate in-package twin of pipekit_jax.JaxModelOp (kept separate so pipekit-train depends only on pipekit core). It can write its weights with serialize_weights, but it can't reload them; pass a JaxModelOp to TrainingLoop for that, and run returns the trained model in the same wrapper.
• 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_ops.Operator.

Public surface used in notebooks and downstream packages. Same constructor signature and _apply shape as pipekit_jax.JaxModelOp. It writes its array leaves with serialize_weights, so a registry can store them and two trained models get different hashes. Swap the import when you need to reload them (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_ops.Operator`.

    Public surface used in notebooks and downstream packages. Same
    constructor signature and ``_apply`` shape as
    ``pipekit_jax.JaxModelOp``. It writes its array leaves with
    ``serialize_weights``, so a registry can store them and two trained
    models get different hashes. Swap the import when you need to reload
    them (``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 serialize_weights(self) -> bytes:
        """Return the module's leaves as ``tree_serialise_leaves`` bytes.

        Unfiltered, so that Python scalar fields are written too.
        ``pipekit_jax.JaxModelOp.with_weights`` reads these bytes against an
        unfiltered skeleton of the same structure.
        """
        buffer = io.BytesIO()
        eqx.tree_serialise_leaves(buffer, self.module)
        return buffer.getvalue()

    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__}

serialize_weights()

Return the module's leaves as tree_serialise_leaves bytes.

Unfiltered, so that Python scalar fields are written too. pipekit_jax.JaxModelOp.with_weights reads these bytes against an unfiltered skeleton of the same structure.

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

def serialize_weights(self) -> bytes:
    """Return the module's leaves as ``tree_serialise_leaves`` bytes.

    Unfiltered, so that Python scalar fields are written too.
    ``pipekit_jax.JaxModelOp.with_weights`` reads these bytes against an
    unfiltered skeleton of the same structure.
    """
    buffer = io.BytesIO()
    eqx.tree_serialise_leaves(buffer, self.module)
    return buffer.getvalue()

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. The adapter synthesises a TrainTask from loop.loss when no explicit task is supplied.

By default evaluation also calls loss_fn, compiled, on ragged batches, and averages its aux metrics per sample across the validation set.

A task may also define eval_fn, which switches evaluation to padded batches. Every batch then has batch_size rows, padded with copies of the last real row, and compiles once. valid is (batch_size,) bool, True for a real row, and the task must leave the other rows out of its metrics. Each batch's metrics are weighted by valid.sum(), so a masked per-sample mean is exact across the set; a metric such as a quantile is not.

Signatures

loss_fn(model, batch, key) -> (scalar_loss, aux_metrics_dict) eval_fn(model, batch, key, valid) -> 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``. The adapter synthesises a `TrainTask` from
    ``loop.loss`` when no explicit task is supplied.

    By default evaluation also calls ``loss_fn``, compiled, on ragged batches,
    and averages its aux metrics per sample across the validation set.

    A task may also define ``eval_fn``, which switches evaluation to padded
    batches. Every batch then has ``batch_size`` rows, padded with copies of the
    last real row, and compiles once. ``valid`` is ``(batch_size,)`` bool, True
    for a real row, and the task must leave the other rows out of its metrics.
    Each batch's metrics are weighted by ``valid.sum()``, so a masked per-sample
    mean is exact across the set; a metric such as a quantile is not.

    Signatures:
        loss_fn(model, batch, key) -> (scalar_loss, aux_metrics_dict)
        eval_fn(model, batch, key, valid) -> 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, final_metrics
tuple[Operator, dict[str, Any]]	(the last training batch), and final_eval_metrics (the validation
tuple[Operator, dict[str, Any]]	set on the final state, empty without a val_dataset) for inclusion
tuple[Operator, dict[str, Any]]	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``, ``final_metrics``
        (the last training batch), and ``final_eval_metrics`` (the validation
        set on the final state, empty without a ``val_dataset``) 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 a directory is set) --------------------
    mngr: ocp.CheckpointManager | None = None
    # A Checkpoint callback sets cadence and retention, and may set the directory.
    checkpoint_cb = next(
        (
            cb
            for cb in loop.callbacks
            if hasattr(cb, "every_n_steps") and hasattr(cb, "keep_last")
        ),
        None,
    )
    checkpoint_dir = _resolve_checkpoint_dir(loop.checkpoint_dir, checkpoint_cb)
    if checkpoint_dir is not None:
        every_n = loop.eval_every_n_steps
        keep_last = 3
        if checkpoint_cb is not None:
            every_n = int(checkpoint_cb.every_n_steps)
            keep_last = int(checkpoint_cb.keep_last)
        mngr = ocp.CheckpointManager(
            directory=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] = {}
    last_eval_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)
                last_eval_metrics = _evaluate(
                    state.model, loop.val_dataset, loop.batch_size, task, eval_key
                )
                _dispatch(loop.callbacks, "on_eval_end", loop, carry, last_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

        # --- Save the final step, which is usually off the cadence ---------
        if (
            mngr is not None
            and train_iter is not None
            and int(state.step) != mngr.latest_step()
        ):
            save_state(
                mngr,
                state,
                int(state.step),
                data_iter_state=train_iter.get_state(),
                force=True,
            )

        # --- 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)
            last_eval_metrics = _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, last_eval_metrics)
    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 = _wrap_trained_model(loop.model_op, 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,
        "final_eval_metrics": last_eval_metrics,
        "sharding": _sharding_info(sharding),
    }
    return trained_model_op, backend_info

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

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
force	bool	Save even if step is off the manager's cadence.	False

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,
    force: bool = False,
) -> 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.
        force: Save even if ``step`` is off the manager's cadence.
    """
    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), force=force)

    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,
    )