Design Decisions¶
Eleven ADRs total. Decisions D1–D7 are about the carrier-agnostic surface. Decisions D8–D11 are about the Equinox adapter specifically (they port and extend the eqx_trainer D1–D6).
When the design lands, each accepted ADR should also be reflected in
a one-paragraph entry in the package CHANGELOG and (where it
materially constrains future work) in AGENTS.md.
D1: Thin orchestration over existing training tools¶
Status: accepted
Context: PyTorch Lightning, Equinox+Optax, and Keras 3 each ship complete, battle-tested training loops with distributed support, mixed precision, profiling, and checkpointing. A fourth implementation cannot improve on any of them.
Decision: pipekit-train wraps these tools through an adapter pattern. It does not reimplement gradient descent, optimiser updates, distributed primitives, or checkpoint serialisation. Each adapter is ~150 LOC of glue.
Consequences: Users debugging a training failure may need to read both pipekit-train and the underlying backend. Mitigation: error messages name the responsible layer; per-backend debugging guides in the docs.
D2: TrainingDataset is a pipekit_ops.Operator¶
Status: accepted
Context: Could be a generator function, an iterator class, or an operator. Generator and iterator forms lose YAML round-trip, content hashing, and composability with the rest of pipekit. Master plan proposes the Operator form.
Decision: TrainingDataset extends pipekit_ops.Operator. It
defines __iter__ returning (x, y) pairs and a content_hash()
returning a stable identifier. Subclasses (CatalogDataset,
SimulationDataset, CachedDataset) inherit operator semantics:
they serialise via get_config / dumps, compose with
Sequential, and integrate with the registry.
Consequences: Iteration semantics need to be defined precisely
for the operator surface — see api/datasets.md for the
"__iter__ is the side-effect entrypoint; _apply is undefined"
rule. Worth it for the composability gain.
D3: TrainingLoop is a StatefulOperator¶
Status: accepted
Context: Master plan and eqx_trainer both treat the training
loop as the headline composable. The pipekit-core StatefulOperator
abstraction (Group M) already exists for cycles and iterative
processes. The training loop's natural shape is exactly
"(carrier_in, state_in) → (carrier_out, state_out)" with state
being (model, optimizer_state, step, epoch, metrics, rng).
Decision: TrainingLoop extends pipekit_ops.StatefulOperator. Its
carry-state is a TrainerCarryState (subclass of
pipekit_ops.CarryState) carrying the six fields above.
TrainingLoop._apply(carrier, state) delegates to the configured
adapter and returns (trained_model_op, new_state).
Consequences: Checkpointing reduces to "serialise the
carry-state". Resume-from-checkpoint is a single
from_dict. Multi-stage training (pretrain → fine-tune) is a
Sequential of two TrainingLoops threading carry-state through.
D4: Loss is a Protocol; common Losses are Operators¶
Status: accepted
Context: eqx_trainer (D4) makes TrainTask a Protocol and
deliberately avoids shipping common losses. The master plan ships
MSE, NLL, KL, Composite as Operators. Both can coexist.
Decision: Ship two surfaces. (1) A runtime-checkable Loss
Protocol (signature: (predicted, target) -> float) for users who
want to plug in their own loss. (2) Common Loss Operators (MSE,
NLL, KL, Composite) that satisfy the Protocol and additionally
provide YAML round-trip + composition. Per-backend, the adapter
either calls the Loss directly or wraps it in the backend's
expected idiom (Equinox: bundles it into a synthesised TrainTask;
Lightning: uses it inside the generated LightningModule.training_step).
Consequences: Two ways to write a loss is an honest cost. The Protocol is the escape hatch; the common Losses are the default.
D5: Adapter pattern, extras-gated¶
Status: accepted
Context: Each backend has a different idiom (Lightning's LightningModule, Equinox's filter_jit, Keras's compile/fit). Heavy backend dependencies (torch, jax, tensorflow) cannot all be hard requirements.
Decision: Each backend lives in one module under
pipekit_train.adapters.<backend>. Each is gated behind a
[<backend>] extra in pyproject.toml. Importing the adapter
module without its extra raises a clean ImportError on first use,
not at import time — so a CLI can branch on what's installed.
Equinox is the v0.1 reference. Lightning and Keras adapters are
scaffolds that raise NotImplementedError in v0.1; targeted v0.2 /
v0.3.
Consequences: Pipekit-train ships with no backend by default.
Users must pip install pipekit-train[equinox] (or another
backend). This matches the pipekit-array and pipekit-experiment
extras model.
D6: Training artifacts integrate with pipekit-experiment¶
Status: accepted
Context: Reproducibility for training runs requires capturing the
training pipeline, dataset content hash, trained model hash, tracker
run id, and registry URI. pipekit-experiment already ships
TrainingArtifact for exactly this purpose.
Decision: TrainingLoop.run() returns
(trained_model_op, TrainingArtifact). The artifact is filled in by
the callback layer: LogToExperiment fills tracker_run_id;
Checkpoint (when configured against a ModelRegistry) fills
model_registry_uri and trained_model_hash; the loop itself fills
training_pipeline_yaml, dataset_hash, backend_info, and
deps_lock. The pipekit-experiment dependency is gated behind the
[experiment] extra.
Consequences: No new reproducibility primitive is invented in
this package. Users without pipekit-experiment get a "naked"
(trained_model_op, dict) return value — the dict is the same
shape as TrainingArtifact.__dict__ but not typed.
D7: Trained models are pipekit Operators¶
Status: accepted
Context: A trained model needs to compose with inference
pipelines without rewrite. pipekit-array.ModelOp and (planned)
pipekit-jax.JaxModelOp already exist as the operator wrappers.
Decision: The Layer 4 adapter's run() returns a
pipekit_ops.Operator. For Equinox, this is the wrapper the user passed in:
a JaxModelOp stays a JaxModelOp, which keeps pipekit-train independent
of pipekit-jax, and an EquinoxModelOp stays one. For Lightning / Keras,
it's a ModelOp.
The wrapper round-trips its weights through a backend-defined
weight blob (bytes) that ModelRegistry.store(..., weights=...)
persists alongside the operator config.
Consequences: A trained emulator drops into
pipekit_cycle.NeuralForward(model_op, dt=...) without any
conversion. A trained classifier drops into any Sequential
preprocessing pipeline. This is the core "train→serve loop is one
operator swap" claim.
D8: Equinox adapter uses Grain + Orbax + Optax¶
Status: accepted (ports eqx_trainer D2, D3)
Context: The eqx_trainer design picks Grain for data, Orbax for checkpointing, and Optax for optimisation. These are the modern JAX-ecosystem choices: deterministic iterators, async multi-device checkpoints, structured gradient transformations.
Decision: The Equinox adapter uses these libraries directly.
grain.DataLoader for the iterator (with
PyGrainCheckpointHandler so preemption recovery restores both
weights and iterator position); orbax-checkpoint.CheckpointManager
with eqx.partition bridge for serialisation; optax for the
optimiser.
Consequences: The [equinox] extra pins six libraries
(equinox, optax, jax, grain, orbax-checkpoint, jaxtyping). All are
single-binary-friendly on CPU+GPU. No TensorFlow dependency leaks
through.
D9: TrainTask Protocol per backend; Callback hooks are optional¶
Status: accepted (ports eqx_trainer D4)
Context: Two related issues. (1) eqx_trainer (D4) makes the
user-supplied training target a TrainTask Protocol (loss_fn +
optional eval_fn). This is the cleanest interface for JAX where
the loss function returns auxiliary metrics from inside the same
eqx.filter_grad. Other backends have their own idioms
(Lightning's LightningModule.training_step). (2) The Callback
surface defines five hooks but most callbacks only care about one
or two; the adapter must tolerate partial implementations.
Decision:
- Each adapter defines its own
TrainTaskProtocol or equivalent. For the Equinox adapter, this is the eqx_trainerTrainTaskProtocol verbatim. For the Lightning adapter (v0.2), the equivalent is aLightningModulefactory function. For Keras (v0.3), it's a(loss, metrics)tuple. The carrier-agnosticLossfrom D4 is the default — if the user doesn't supply a backendTrainTask, the adapter synthesises one from theLossOperator. - The
CallbackProtocol is not@runtime_checkable.runtime_checkablewould require every callback to implement all five hooks (otherwiseisinstance(cb, Callback)returns False), which contradicts the partial-implementation model. The Protocol stays as the static-typing contract for the full hook surface; runtime dispatch in adapters usesgetattr(cb, hook, None).LossandMetricWriterare runtime-checkable because every implementation has the same single-method surface.
Consequences: Power users see backend idioms; default users
only see Loss. The synthesised path is what the worked examples
in examples/ use. Partial Callback implementations (e.g. a
metric logger that only cares about on_step_end) are valid by
design, and the test suite encodes this contract structurally
rather than via isinstance.
D10: Per-step is the unit of training, not per-epoch¶
Status: accepted
Context: Master plan uses n_epochs. eqx_trainer uses
num_steps. Grain iterators are typically infinite (num_epochs=None).
Per-epoch semantics break down when the dataset is a streaming
simulator or when steps_per_epoch isn't well-defined.
Decision: The training loop's unit is the step. TrainingLoop
takes max_steps (required); steps_per_epoch is optional and
purely cosmetic (controls when on_epoch_end callbacks fire).
Eval and checkpoint cadences are every_n_steps. Datasets are
expected to be infinite iterators; finite datasets are wrapped to
repeat.
Consequences: Existing-user mental model of "10 epochs" needs translation to "10 × steps_per_epoch steps". Documented prominently in the examples.
D11: TrainerConfig is the operator's config dict, not a separate dataclass¶
Status: accepted (refines eqx_trainer D5)
Context: eqx_trainer (D5) ships a TrainerConfig dataclass and
explicitly avoids YAML / ml_collections. The pipekit world already
has a YAML round-trip via pipekit_ops.serial.dumps. A separate config
dataclass would duplicate machinery.
Decision: TrainingLoop is constructed with the same flat-config
shape that the eqx_trainer TrainerConfig proposes (max_steps,
seed, checkpoint_every_n_steps, max_checkpoints_to_keep,
eval_every_n_steps, log_every_n_steps, …), but those fields live
directly on the operator and are surfaced through its
get_config(). The YAML round-trip is free. Users who prefer the
dataclass-only style can ignore the YAML side and treat
TrainingLoop exactly like the eqx_trainer's Trainer +
TrainerConfig pair.
Consequences: Slight tension with eqx_trainer D5's "config files are the caller's responsibility". Resolved by saying: pipekit-train provides the YAML facility (via the operator surface) but does not require it. Callers using a dataclass-style API see no YAML.
D12: Sharding is an optional ShardingSpec, not a new code path¶
Status: accepted (resolves boundaries.md Q1, issue #15)
Context: The v0.1 Equinox adapter was single-host, single-device. Multi-GPU (model parallelism, larger batches) and multi-host (training across nodes) are the most-requested extension. The risk was forking the adapter into a parallel "distributed" code path.
Decision: Add one optional, frozen ShardingSpec(mesh, model_pspec,
data_pspec) and thread it through the existing path rather than
branching:
TrainState.create(..., sharding=spec)places model + optimiser array leaves onspec.model_sharding; the step counter is replicated.train_stepis unchanged —eqx.filter_jitcompiles sharding-aware automatically once its inputs are sharded (replicated model + a batch sharded alongdata_pspecgives data-parallel SPMD with XLA inserting the collectives). No explicitpjitwrapper is needed for the data-parallel case.- The data iterator places each batch on
spec.data_sharding; for multi-host it first reads a disjoint process-shard of thegrainMapDataset(map_ds[process_index::process_count], the modern equivalent ofgrain.ShardByJaxProcess) and assembles the global array withjax.make_array_from_process_local_data. - Orbax restore is sharding-aware for free: the restore template is the
already-sharded
TrainState, soto_shape_dtype_structcarries each leaf's sharding intoStandardRestore. TrainingLoop.shardingis typedAny | Noneso the carrier-agnostic core never imports JAX (the spec lives inadapters.equinox).
Consequences: sharding=None is byte-for-byte the old single-device
path. The default (data-parallel) spec is model_pspec=PartitionSpec()
(replicated) + data_pspec=PartitionSpec("data"); uniform model-parallel
specs work when leaf ranks are compatible. Verified on a CPU-simulated
2-device mesh and a scripted 2-process jax.distributed smoke. The
Lightning adapter keeps its own native distributed support (D5).
D13: NumPyro = two task-first adapters — numpyro-svi + numpyro-mcmc¶
Status: accepted (design; implementation planned)
Context: A NumPyro backend would make Bayesian inference a first-class
TrainingLoop target, complementing the pipekit-evaluate benchmark ladder
(MCMC = rung-2 oracle, SVI = rung-4). But NumPyro's "training" is inference
over a probabilistic model — it does not match the carrier-agnostic
Loss(pred, target) surface (D4) — and its two paradigms have very
different loop shapes: SVI's svi.update is a per-step gradient step,
while MCMC's mcmc.run is a single blocking sampler.
Decision:
- Two adapters, split by paradigm, not one method-dispatched module:
backend="numpyro-svi"andbackend="numpyro-mcmc". SVI genuinely is the per-stepTrainingLoop(max_steps,optimizer_config, per-step eval/checkpoint,_BatchSourceminibatching all apply;SVI(optim=...)takes an optax transformation). MCMC is a singlemcmc.runover the full dataset (begin/end callbacks only; those loop fields do not apply). Splitting removes the "ignored for nuts" caveats and lines NumPyro up with BlackJAX (D14):numpyro-mcmcandblackjaxare sibling sampler backends,numpyro-sviis the variational one. The two NumPyro backends share the same[numpyro]extra — a deliberate, minor exception to D5's "one extra per backend" (the dependency is identical; only the paradigm differs). - Task-first, shared task. Both require
loop.task(one sharedNumpyroTask: model + predictive config + SVI fields + MCMC fields), notloop.loss— the objective is the model's ELBO / joint log-density. Nomethodfield: the backend string selects the paradigm, so the same task works under either backend. Uses the D9 per-backendTrainTaskseam; passing onlyloss=errors (no loss→task synthesis). - Shared seam.
NumpyroPredictiveOp, the model→logdensity bridge, and thePredictivewrapping live inpipekit_train.adapters.bayes, reused by both NumPyro adapters and the BlackJAX adapter (D14). The artifact is a posterior-predictiveOperatorreturning the mean of a configuredpredictive_site; weight-blob = the params / posterior-samples PyTree (D7).
Consequences: Two thin adapters over one shared seam, rather than one
overloaded module. No hard dependency (gated behind [numpyro]). Full design
in numpyro_adapter.md; concise entries in
api/adapters.md. Scope: SVI + NUTS.
D14: BlackJAX is a separate backend, sharing a Bayesian seam with NumPyro¶
Status: accepted (design; implementation planned)
Context: BlackJAX is also wanted as a backend. But BlackJAX is a
sampler library, not a PPL: it operates on a logdensity_fn and is
PPL-agnostic (it consumes log-densities from NumPyro, PyMC, or hand-written
functions). The question was whether to bundle it into the NumPyro adapter
or ship it separately.
Decision: A separate backend="blackjax" module, gated behind a
[blackjax] extra — not a NumPyro mode — for four reasons:
- Different tool / extra / contract (D5). The
[numpyro]extra must not pull BlackJAX, and BlackJAX must not require NumPyro (that would deny its PPL-agnostic purpose). Its task is alogdensity_fn, not a model. - It fits the per-step loop better than NumPyro. BlackJAX's
kernel.step(rng, state) -> (state, info)is one sample per step, so the adapter reuses the per-stepTrainingLoop(per-sample diagnostics, checkpoint, early-stop) — unlike NumPyro's blockingmcmc.run, which the NumPyro adapter runs as a singlefit(D13). - A larger algorithm zoo — MCLMC, stochastic-gradient MCMC
(
sgld/sghmc, minibatch MCMC via_BatchSource), tempered SMC,pathfinder/svgd. - Cross-checking. A second, independent engine for the same model strengthens the benchmark-ladder oracle story.
Shared seam, not bundling. Both adapters emit the same
posterior-predictive Operator (D7) and both can start from a NumPyro
model. The model→logdensity bridge + Predictive wrapping is factored into
pipekit_train.adapters.bayes; a NumpyroTask is one accepted task source
for the BlackJAX adapter. Net UX: write a NumPyro model once, switch the
MCMC engine by flipping backend="numpyro-mcmc" ↔ "blackjax".
Consequences: Two thin Bayesian adapters over one shared seam rather
than one overloaded module. Full design (user story, math, CS background,
API, examples, build order, references) in
blackjax_adapter.md. Scope agreed as NUTS + SG-MCMC
first, MCLMC / SMC / VI as follow-ons.
D15: Drop the Lightning backend¶
Status: accepted (supersedes the Lightning part of D5)
Context: D5 listed Lightning as a scaffolded backend for v0.2. The project does not want to maintain PyTorch training integration.
Decision: pipekit-train does not support Lightning or PyTorch
training. The lightning backend stub, the [lightning] extra and
the Lightning documentation are removed. Keras (v0.3 scaffold) and the
Equinox, NumPyro and BlackJAX backends are unaffected.
Consequences: backend="lightning" is no longer a valid choice.
Whether pipekit-array keeps torch carrier support is a separate
question, not decided here.
D16: Padded evaluation is opt-in through TrainTask.eval_fn¶
Status: accepted
Context: Evaluation compiles loss_fn on ragged batches, which gives exact
per-sample means and no API change. A ragged batch can't split evenly across
devices, though, and a task can't tell padding from data unless something passes
it a mask.
Decision: A task may define eval_fn(model, batch, key, valid) -> metrics.
Its presence, checked with hasattr, switches the Equinox adapter to padded
batches. Padding copies the last real row across the batch pytree, so padded rows
are finite and in distribution. Each batch's metrics are weighted by valid.sum().
Consequences: A masked per-sample mean is exact across the validation set, and
the step compiles once. A task that doesn't define eval_fn is unaffected. A
non-mean metric such as a quantile isn't exact across batches, because returning
sufficient statistics would make every task's contract heavier. Tasks can't ignore
the mask by accident, because valid is a required argument of eval_fn, whereas
adding it to loss_fn would let a task silently average padded rows. The adapter
doesn't shard evaluation yet; this makes that possible.
Resolved Questions¶
| Question | Resolution |
|---|---|
| Backend strategy | Adapter pattern, extras-gated (D5) |
| Carrier-agnostic surface | Datasets + Losses as Operators (D2, D4) |
| Loop primitive | StatefulOperator (D3) |
| Reproducibility | pipekit-experiment.TrainingArtifact (D6) |
| Train→serve composition | Trained models are Operators (D7) |
| JAX stack | Grain + Orbax + Optax (D8) |
| User loss interface | Per-backend TrainTask Protocol (D9) |
| Training unit | Per-step (D10) |
| Config representation | Operator.get_config() — YAML for free (D11) |
| Distributed training | Delegated to backend (JAX sharding for Equinox) |
| — see boundaries.md | |
| Hyperparameter search | Deferred to v0.2 sweep.py — see boundaries.md |
Open Questions (see boundaries.md)¶
- Q1: Multi-device sharding for the Equinox adapter.
- Q2: Gradient accumulation (Optax MultiSteps vs trainer-level).
- Q3: Mixed precision (per-backend, no carrier-agnostic API).
- Q4: Stateful layers (
eqx.nn.Statefor BatchNorm-like layers). - Q5: Streaming / online updates as a v0.3 follow-up.