Skip to content

Layer 4 — Backend Adapters

One module per backend under pipekit_train.adapters.*. Each exports a single run(loop) → (trained_model_op, backend_info) function. v0.1 ships the Equinox adapter as the reference; Keras is scaffolded.

Equinox (v0.1 reference)

pipekit_train.adapters.equinox. Built on the eqx_trainer design: Grain DataLoader, Orbax CheckpointManager, Optax optimiser, eqx.filter_jit train step, eqx.partition checkpoint bridge.

Extras: pipekit-train[equinox].

Module surface

# pipekit_train/adapters/equinox.py


class TrainState(eqx.Module):
    """Equinox-side training state. Stored inside TrainerCarryState.opt_state.

    A nested-state design: TrainerCarryState (the pipekit-side carry
    that everyone sees) holds an opaque opt_state field; for the
    Equinox adapter, that field is itself a TrainState. The split
    keeps the JAX-specific PyTree concerns inside the adapter and
    the carrier-agnostic concerns inside the loop.
    """

    model: eqx.Module
    opt_state: optax.OptState
    step: int

    @staticmethod
    def create(
        model: eqx.Module, optimizer: optax.GradientTransformation
    ) -> "TrainState":
        opt_state = optimizer.init(eqx.filter(model, eqx.is_array))
        return TrainState(model=model, opt_state=opt_state, step=0)


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

    Provide `loss_fn`. The adapter synthesises a TrainTask from
    `loop.loss` if none is supplied. Evaluation compiles `loss_fn` and
    averages its aux metrics per sample, unless the task defines the
    optional `eval_fn(model, batch, key, valid) -> metrics_dict`.
    """

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


@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 (see api/primitives.md)."""


def save_state(mngr: ocp.CheckpointManager, state: TrainState, step: int) -> None: ...
def restore_state(
    mngr: ocp.CheckpointManager, template: TrainState, step: int
) -> TrainState: ...


def run(loop: TrainingLoop) -> tuple[Operator, dict[str, Any]]:
    """Train `loop` end-to-end with the Equinox+Optax+Grain+Orbax stack.

    Returns:
        trained_model_op: A pipekit_jax.JaxModelOp (or domain
            subclass) wrapping the trained eqx.Module.
        backend_info: Dict with keys: backend ("equinox"), jax_version,
            equinox_version, optax_version, devices, total_seconds.
    """

Padded evaluation with eval_fn

Defining eval_fn switches evaluation to padded batches, where every batch has batch_size rows and compiles once. Padded rows are copies of the last real row, and valid is True for real rows.

class Regression:
    def loss_fn(self, model, batch, key):
        x, y = batch
        err = (jax.vmap(model)(x) - y) ** 2
        return err.mean(), {"mse": err.mean()}

    def eval_fn(self, model, batch, key, valid):
        x, y = batch
        err = ((jax.vmap(model)(x) - y) ** 2).mean(axis=-1)  # (batch_size,)
        return {"mse": jnp.where(valid, err, 0.0).sum() / valid.sum()}

Each batch's metrics are weighted by valid.sum(), so a masked per-sample mean is exact across the validation set. A metric that isn't a mean of per-sample terms, such as a quantile, is only exact when the whole set fits in one batch.

What run does

  1. Resolve the optimizer. optimizer_config = {"name": "adamw", "lr": 3e-4, ...} → optax.chain(optax.adamw(**kwargs), ...). Schedule support: if lr is a dict like {"schedule": "cosine_decay", "init_value": 3e-4, "decay_steps": loop.max_steps}, build an optax.cosine_decay_schedule and nest the optimiser inside optax.inject_hyperparams.
  2. Resolve the task. If loop.task is set, use it. Otherwise synthesise: class SynthTask: def loss_fn(self, model, batch, key): pred = jax.vmap(model)(batch["x"]); l = loop.loss(pred, batch["y"]); return l, {"loss": l}.
  3. Build the data iterator. Wrap loop.dataset as a grain.MapDataset source if it isn't already one; add a grain.Batch(loop.batch_size) operation; wrap in a grain.DataLoader(worker_count=2, ...). Same for loop.val_dataset.
  4. Set up the checkpoint manager. If loop.checkpoint_dir, build an ocp.CheckpointManager with save_interval_steps=loop.callbacks[Checkpoint].every_n_steps and max_to_keep=loop.callbacks[Checkpoint].keep_last, using a CompositeCheckpointHandler so the Grain iterator state checkpoints alongside the model.
  5. Initialise TrainState and TrainerCarryState. Pull the model from loop.model_op (it must be a JaxModelOp or a subclass exposing .module returning the bare eqx.Module). Build the carry-state with the right RNG split.
  6. Resume. If checkpoint manager has a latest_step, restore TrainState + Grain iterator state.
  7. Run the per-step loop. For each step up to loop.max_steps: call train_step; periodically log (log_every_n_steps), evaluate (eval_every_n_steps), checkpoint, run callbacks.
  8. Build the trained model_op. Wrap the final eqx.Module back into a JaxModelOp (or the original model_op's class) and return it together with backend_info.

Sharding (multi-device / multi-host) — v0.2

Pass an optional ShardingSpec to enable data-parallel (and uniform model-parallel) training (issue #15, ADR D12):

import jax
from jax.sharding import Mesh, PartitionSpec
from pipekit_train import TrainingLoop
from pipekit_train.adapters.equinox import ShardingSpec

mesh = Mesh(jax.devices(), axis_names=("data",))
spec = ShardingSpec(
    mesh=mesh,
    model_pspec=PartitionSpec(),  # replicate the model on every device
    data_pspec=PartitionSpec("data"),  # shard each batch along its leading axis
)

loop = TrainingLoop(
    model_op=model_op,
    dataset=dataset,
    loss=loss,
    batch_size=256,  # must be divisible by the "data" axis size
    sharding=spec,  # None => single-device (the default)
)
trained_op, artifact = loop.run()

How it threads through run:

  • TrainState — model + optimiser array leaves are placed on spec.model_sharding; the step counter is replicated.
  • train_step — unchanged; eqx.filter_jit compiles sharding-aware once its inputs are sharded (replicated model + a data-sharded batch give data-parallel SPMD, XLA inserting the collectives).
  • Data iterator — each batch is placed on spec.data_sharding. For multi-host (an explicit spec and process_count > 1), each process reads a disjoint shard of the grain MapDataset (map_ds[process_index::process_count], the modern equivalent of grain.ShardByJaxProcess) and batches the local size batch_size // process_count, so the per-process batches assemble via jax.make_array_from_process_local_data into a global batch of exactly batch_size. Without a spec the single-device path is preserved even under a distributed-initialised JAX. (batch_size must divide process_count; multi-host needs an indexable dataset, not a streaming one.)
  • Checkpointing — Orbax restore is sharding-aware automatically: the restore template is the already-sharded TrainState.
  • backend_info["sharding"] records the mesh shape and partition specs for the artifact.

Verified on a CPU-simulated 2-device mesh (tests/adapters/ test_equinox_sharding.py) and a scripted 2-process jax.distributed smoke (scripts/multihost_sharding_smoke.py).

Notes

  • sharding=None (the default) is the original single-host, single-device path, unchanged.
  • The adapter does not require the user to install pipekit-jax. If it isn't installed, the adapter returns a minimal JaxModelOp-like wrapper defined inline (the wrapper just calls the eqx.Module's __call__). The full pipekit-jax.JaxModelOp adds shape inference and weight-blob round-trip.
  • The Grain iterator is the only JAX-side data-loading dependency. PyTorch DataLoader is explicitly not used.

Keras (v0.3 scaffold)

pipekit_train.adapters.keras. Extras: pipekit-train[keras].

v0.1 ships with the same NotImplementedError stub.

Planned design

When v0.3 lands, the adapter will:

  1. Translate loop.loss into a keras.losses.Loss subclass (or accept the user's pre-built tf.keras.losses instance via loop.task).
  2. Configure the optimiser via keras.optimizers.get using optimizer_config.
  3. Wrap loop.dataset as a tf.data.Dataset.from_generator(...).
  4. Call model.compile(...) then model.fit(...).
  5. Wrap the trained keras.Model in pipekit-array.ModelOp and return.

Keras 3's multi-backend (TF / JAX / torch) means the Keras adapter gives a second JAX path orthogonal to the Equinox one — useful for projects that already standardise on Keras layers.

NumPyro (implemented) — two Bayesian backends: numpyro-svi + numpyro-mcmc

pipekit_train.adapters.numpyro_svi and …numpyro_mcmc. Extras: pipekit-train[numpyro] (numpyro>=0.15, optax, grain; JAX via NumPyro). See ADR D13 and the full design in ../numpyro_adapter.md.

"Training" here is Bayesian inference. NumPyro is split into two backends by paradigm (the per-step vs blocking loop shapes are too different to method-dispatch):

  • numpyro-svi (variational, per-step) — svi.update returns (state, loss) per step, so it is the per-step loop: callbacks / metric_writer / eval / checkpoint / max_steps / _BatchSource minibatching all apply, and SVI(optim=...) takes an optax transformation so optimizer_config is reused.
  • numpyro-mcmc (sampling, single call) — one mcmc.run (NUTS) over the full dataset (begin/end callbacks only; the per-step fields don't apply).

Both are task-first (D9): loop.task is one shared NumpyroTask (no method field — the backend selects the paradigm, so the same task works under either); the carrier-agnostic Loss (D4) does not apply. The artifact is a NumpyroPredictiveOp (posterior-predictive Operator) from the shared adapters.bayes seam (D7). They fill the benchmark ladder's rung-2 (numpyro-mcmc) and rung-4 (numpyro-svi); numpyro-mcmc is a sibling of the BlackJAX sampler backend (below). Full design — user story, math, CS background, API, examples, build order, references: ../numpyro_adapter.md.

Unsupervised fits, kernel settings and full posterior draws:

def model(data=None, *, prior_scale=10.0):
    mu = numpyro.sample("mu", dist.Normal(0.0, prior_scale))
    sigma = numpyro.sample("sigma", dist.HalfNormal(5.0))
    with numpyro.plate("n", 32 if data is None else data.shape[0]):
        numpyro.sample("obs", dist.Normal(mu, sigma), obs=data)


task = NumpyroTask(
    model,
    supervised=False,  # a dataset of single items
    model_kwargs={"prior_scale": 10.0},
    kernel_kwargs={"target_accept_prob": 0.995},  # NUTS(model, **kernel_kwargs)
)
op, artifact = TrainingLoop(..., task=task, backend="numpyro-mcmc").run()

op.posterior_samples()["mu"]  # (num_samples,)
op.predictive_draws(None)["obs"]  # (num_samples, 32): None draws new data
op(None)  # the posterior-predictive mean, as before

# after Checkpoint(registry=registry); the model is a closure, so it isn't stored
op = NumpyroPredictiveOp.from_registry(
    registry, artifact.trained_model_hash, model=model
)

BlackJAX (planned) — sampler-library backend

pipekit_train.adapters.blackjax. Extras: pipekit-train[blackjax] (blackjax>=1.2, optax, grain; numpyro only for the NumPyro-model task path). See ADR D14 and the full design in ../blackjax_adapter.md.

Separate from NumPyro, by design. BlackJAX is a sampler library, not a PPL: it operates on a logdensity_fn, and is PPL-agnostic. So it is its own backend (different tool / extra / contract) — but it interoperates: a BlackjaxTask accepts a raw logdensity_fn or a NumpyroTask (bridged via numpyro.infer.util.initialize_model), and both adapters share the model→logdensity + Predictive seam (adapters.bayes). Net UX: write a NumPyro model once, flip backend="numpyro-mcmc" ↔ "blackjax".

Why a separate adapter pays off: BlackJAX's kernel.step(rng, state) -> (state, info) is one sample per step, so it reuses the per-step loop (per-sample diagnostics / checkpoint / early-stop) — a cleaner fit than NumPyro's blocking mcmc.run. It also unlocks an algorithm zoo NumPyro lacks: MCLMC, stochastic-gradient MCMC (sgld/sghmc — minibatch MCMC via _BatchSource), tempered SMC, and pathfinder/svgd VI. Full design: ../blackjax_adapter.md.

Adapter selection

TrainingLoop.run() selects the adapter by string:

loop = TrainingLoop(..., backend="equinox")  # uses adapters.equinox
loop = TrainingLoop(..., backend="keras")  # uses adapters.keras
loop = TrainingLoop(..., backend="numpyro-svi")  # adapters.numpyro_svi  (planned)
loop = TrainingLoop(..., backend="numpyro-mcmc")  # adapters.numpyro_mcmc (planned)
loop = TrainingLoop(..., backend="blackjax")  # adapters.blackjax     (planned)

The selection happens at run() time, not at construction, so a TrainingLoop YAML can be authored with backend="keras" even on a machine that hasn't installed [keras] — the missing extra is reported at run time only.