pipekit-train

Summary¶

Decision. pipekit-train stays a thin layer over existing training tools, as both the scoping report and ADR D1 say, and we keep its 0.0.2 design: step-based loops (D10), run() returning (model_op, TrainingArtifact) (D6), Equinox with Grain and Orbax (D8), and the Bayesian backends (D13, D14). Five workstreams finish it.

• A. Datasets. CatalogDataset is missing (A1). SimulationDataset crashes on its cycle= path and hides prior errors (A2), and its content hash ignores the simulator, so a cache can serve stale data (A3). CachedDataset reads one zarr item at 8.2 ms and only supports zarr (A4). XarrayWindowDataset imports a geopatcher module that no longer exists (A5).

• B. The trained model’s handoff. The Equinox backend returns an EquinoxModelOp which cannot be stored with its weights or reloaded from LocalModelRegistry, whereas wrapping it in pipekit-jax’s JaxModelOp works (B1). We also fill the artifact’s model hash without requiring a registry (B2).

• C. Loop components. Checkpoint ignores save_dir and does not save the final step (C1), HyperSweep ranks trials on the last training batch instead of on validation (C2), ValidationStep is never used by the loop (C3), and NumPyro SVI runs an un-jitted update at 76 ms per step (C4).

• D. Amortized inference. The data path and NLL work, but every user writes their own distribution head and dataset flip. So we ship a recipe and a Gaussian head, and we decide where conditional density estimators live (D1).

• E. The other two backends. We build the Lightning (E1) and Keras 3 (E2) adapters that ADR D5 schedules for v0.2 and v0.3.

End state. The scoping report’s three training shapes all run from pipekit-train: 1) supervised training from a catalog, 2) emulator training from a pipekit-cycle forward model, and 3) amortized inference. Each one produces an operator that drops into an inference pipeline or a cycle, reloads from the registry with its weights, and carries an artifact whose hashes identify its data and its weights.

Ownership rules¶

pipekit-train owns the glue between the data, a model and a training tool, whereas the tools own the training maths and the model libraries own the models (Figure 1).

pipekit-train owns the datasets, losses, the loop, callbacks, sweeps, writers and the adapters to training tools. Its core only imports pipekit. JAX, torch and Keras arrive through extras, and pipekit_experiment is only imported by LogToExperiment and the artifact. It never reimplements an optimizer step or a sampler, it never defines its own artifact type, and it has no model-definition DSL (see its boundaries.md).

The neighbours (the training tools, pipekit-cycle, pipekit-experiment and the model libraries) follow the stack rules.

What ships¶

Table 1 lists the public surface by layer.

The package’s tests pass: 228 tests in 76 s with geopatcher shimmed (see stage 8), plus the BlackJAX (9 of 9) and NumPyro suites, which are marked slow.

The running example: an emulator, a retrieval and a stored archive¶

The stages reuse the two problems of Report 10, so that we can check a trained model against the exact operator it replaces. In the six-step cycle, stage 2 is step 3 of a forecast (and stage 3 runs that emulator forward), and stages 6 and 7 are steps 5 and 2 of a retrieval. Report 17 runs all six steps on the same retrieval.

The transport model is the upwind advection of Report 10 on a periodic 40-cell transect, $\vect{x}_{t+1} = \mat{F} \vect{x}_t$ with $\mat{F} = 0.6 \mat{I} + 0.4 \mat{S}$. Here $\vect{x}_t \in \mathbb{R}^{D_x}$ is the state at step $t$ with $D_x = 40$, $\mat{F} \in \mathbb{R}^{D_x \times D_x}$ is the linear transport model, $\mat{S} \in \mathbb{R}^{D_x \times D_x}$ is the periodic one-cell shift and $\mat{I}$ is the identity (the code keeps the array name M for $\mat{F}$). An emulator is a learned $\hat{\mat{F}} \in \mathbb{R}^{D_x \times D_x}$, and the target is $\hat{\mat{F}} = \mat{F}$.

The retrieval uses the five-band sensor of Report 17 unchanged, $y_b = e^{-k_b \alpha} + \varepsilon_b$ for the bands $b = 1, \ldots, 5$, with the absorption coefficients $\vect{k} = (0.002, \ldots, 0.010) \in \mathbb{R}^5$, the noise standard deviation $\sigma_\varepsilon = 10^{-3}$ and the column enhancement $\alpha \in [0, 300]$, a scalar. We stack the bands into the spectrum $\vect{y} \in \mathbb{R}^5$ and write the noise-free spectrum as $h(\alpha) \in \mathbb{R}^5$. An amortized retrieval learns $p(\alpha \mid \vect{y})$ once, for every $\vect{y}$, whereas a per-scene retrieval (Report 10’s Gauss–Newton, or MCMC) solves for one $\vect{y}$ at a time.

The archive. Fifty hourly states of the transect, written to a Zarr store, stand in for the stored model states of Report 13.

Stage 1 · Training data from a forward model (pipekit-cycle → pipekit-train)¶

TL;DR. We turn the transport model of Report 10 into 2,000 input-output pairs with fixed seeds and with train, validation and test splits.

Problem. A SimulationDataset draws each input $\vect{x}_n \in \mathbb{R}^{D_x}$ from a prior $p(\vect{x})$ with the seed $s_0 + o + n$, where $s_0$ is the base seed, $o$ is the split’s offset and $n$ indexes the samples, and it returns the pair $(\vect{x}_n, \mat{F} \vect{x}_n)$. The splits are 80, 10 and 10 % of the seed range, so they never share a draw.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# Shapes: D_x = 40 (the code's N); N = 2,000 pairs; train 1,600 · val 200 · test 200
import pipekit_cycle as pc
import pipekit_train as pt

fwd = pc.CallableForward(lambda x, dt: M @ x, dt=1.0)    # F x: (D_x,) → (D_x,)    today
class WhitePrior(pk.Operator):                           # x ~ 𝒩(0, I)    user
    def _apply(self, seed=0):
        return np.random.default_rng(seed).normal(size=N)

train = pt.SimulationDataset(fwd, WhitePrior(), n_samples=2000, seed=0)  # today
val = train.with_split("val")                          # disjoint seeds   today
len(train), len(val)                                   # → 1600, 200      today
train.content_hash()                                   # sha256 of the recipe

pt.SimulationDataset(fwd, WhitePrior(), 2000, cycle=pc.Cycle(fwd, 24))[0]
# → TypeError: Cycle._apply() missing 1 required positional argument: 'state'

The one-step path works. However, the trajectory path, cycle=, calls self.cycle(params) without a carry state, so every pipekit-cycle cycle raises (dataset.py:284). A second defect sits just above it. The prior is called as prior(seed=...) inside try … except TypeError, and any TypeError raised inside the prior silently falls back to prior() with no seed, which means every sample comes out identical. Workstream A2 fixes both.

Stage 2 · Train the emulator (pipekit-train → Equinox)¶

TL;DR. We fit a linear emulator $\hat{\mat{F}}$ by least squares on the pairs and see how the choice of prior decides what it learns.

Problem. We take $N$ pairs $(\vect{x}_n, \mat{F} \vect{x}_n)$ and a linear model $\hat{\mat{F}} = \mat{W}$, whose parameters are the entries of the weight matrix $\mat{W} \in \mathbb{R}^{D_x \times D_x}$, so $\gvect{\theta} = \op{vec}(\mat{W}) \in \mathbb{R}^{D_x^2}$. The mean-squared error (1),

is zero at $\mat{W} = \mat{F}$, and the minimiser is only unique when the inputs span $\mathbb{R}^{D_x}$, that is when the stacked inputs $\mat{X} \in \mathbb{R}^{N \times D_x}$ (one $\vect{x}_n^\top$ per row) have full rank. The loop minimises it with Adam at a step size $\gamma = 10^{-2}$ on batches $\mathcal{D}^{(i)}$ of $N_\text{batch} = 64$ pairs, where $i$ counts the optimiser iterations.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# Shapes: W (D_x, D_x); batch 64; 25 steps per epoch
import equinox as eqx
from pipekit_train.adapters.equinox import EquinoxModelOp

loop = pt.TrainingLoop(
    model_op=EquinoxModelOp(eqx.nn.Linear(N, N, use_bias=False, key=key)),
    dataset=train, val_dataset=val, loss=pt.MSE(),
    optimizer_config={"name": "adam", "lr": 1e-2},
    max_steps=6000, steps_per_epoch=25, batch_size=64, eval_every_n_steps=50,
    callbacks=(pt.EarlyStopping(metric="mse", patience=5),),
    seed=0,
)                                                                       # today
trained, artifact = loop.run()            # (EquinoxModelOp, TrainingArtifact)
# white prior:  stops at step 500;  val MSE 5.3e-14;  max |W − F| = 3.8e-7
# smooth prior: stops at step 2200; val MSE 3.4e-5;   max |W − F| = 0.343

Both runs fit their data (Table 2), where $\sigma_{\min} / \sigma_{\max}$ is the ratio of the smallest to the largest singular value of $\mat{X}$. However, the inputs of the smooth prior nearly lie in a low-dimensional subspace, so $\mat{W}$ is only fixed on that subspace and is arbitrary off it. In other words, a low training loss is not a correct emulator. This is a property of the data and not of the code, and it is the reason a SimulationDataset’s prior is part of its hash.

Stage 3 · Drop the emulator into a cycle (pipekit-train → pipekit-cycle)¶

TL;DR. We replace the forward model of a six-hour forecast with the emulator, and we test it on states that the prior did and did not cover.

Problem. The emulated forecast is $\hat{\vect{x}}_6 = \mat{W}^6 \vect{x}_0$, where $\vect{x}_0 \in \mathbb{R}^{D_x}$ is the initial state. Its error $\Vert \mat{W}^6 \vect{x}_0 - \mat{F}^6 \vect{x}_0 \Vert_\infty$ is only small when $\vect{x}_0$ lies where $\mat{W}$ was fixed.

1
2
3
4
5
# Shapes: x₀ (D_x,)
emulator = pc.NeuralForward(trained, dt=1.0)                       # today
x6, _ = pc.Cycle(emulator, n_steps=6)(x0, None)                    # today
# white-prior W:  max |x̂₆ − F⁶ x₀| = 1.1e-6 (smooth x₀), 1.3e-6 (white x₀)
# smooth-prior W: max |x̂₆ − F⁶ x₀| = 0.017 (smooth x₀), 0.533 (white x₀)

The handoff that the scoping report promised works as written, since NeuralForward wraps the trained EquinoxModelOp directly and the cycles of Report 10 accept it. The smooth-prior emulator is accurate on smooth states but wrong by half the signal on white ones. So an emulator is only as general as its prior.

Stage 4 · Callbacks, checkpoints and the artifact (pipekit-train)¶

TL;DR. We watch the smooth-prior run through its callbacks and see what it records about itself.

Problem. A training run must be stoppable, resumable and traceable, which asks for an early-stopping rule, periodic checkpoints, a metric log, a tracker, and an artifact that names the data and the code.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# Shapes: as stage 2
loop = pt.TrainingLoop(
    ..., callbacks=(
        pt.EarlyStopping(metric="mse", patience=5, min_delta=1e-9),
        pt.Checkpoint(every_n_steps=250, keep_last=2, save_dir="unused",
                      registry=px.LocalModelRegistry("registry")),
        pt.LogToExperiment(tracker, run_name="emu"),
    ),
    metric_writer=pt.JSONLWriter("metrics.jsonl"),
    checkpoint_dir="ckpt",
)                                                               # today
trained, art = loop.run()
# early stop at step 2200; ckpt/ holds 1750 and 2000 (not 2200)
# metrics.jsonl: 44 lines; tracker: start, step and 44 eval/* events, end
art.dataset_hash, art.model_registry_uri          # set               today
art.backend_info["final_step"], art.deps_lock     # 2200, uv pip freeze
loop.run()                                        # resumes at step 2000

Early stopping, the metric log, the tracker and auto-resume all work. However, Checkpoint has two defects. It writes to the loop’s checkpoint_dir and ignores its own save_dir (the adapter only reads every_n_steps and keep_last). In addition, the final step is only saved if it is a multiple of every_n_steps, so a run that stops early loses its last 200 steps on resume. Workstream C1 fixes both.

Stage 5 · Store and reload the emulator (pipekit-train → pipekit-experiment)¶

TL;DR. We put the trained emulator in the registry and load it back with its weights.

Problem. A registered model must reload to the same outputs, so after a store and a load, $\max \vert \mat{W}_\text{loaded} - \mat{F} \vert$ should equal the trained $3.8 \times 10^{-7}$.

1
2
3
4
5
6
7
8
9
10
11
12
# Shapes: W (D_x, D_x)
reg = px.LocalModelRegistry("registry")
# Today, through Checkpoint(registry=reg):
reg.load_weights(art.trained_model_hash)          # → None: no weights.bin
reg.load(art.trained_model_hash)
# → RuntimeError: from_state cannot reconstruct EquinoxModelOp … forbid_in_yaml

# What works today: re-wrap as pipekit-jax's JaxModelOp
j = JaxModelOp(trained.module)                                      # today
h = reg.store(j, weights=j.serialize_weights())                     # today
loaded = skeleton.from_registry(reg, h)   # skeleton: a JaxModelOp   today
# → max |W_loaded − F| = 3.8e-7; a 6-step NeuralForward cycle agrees too

The loop always returns an EquinoxModelOp, even when it was given a JaxModelOp, and EquinoxModelOp has no serialize_weights. So Checkpoint’s registry.store writes the operator’s config and metadata but no weights. The design docs (ADR D7) say that the Equinox adapter returns a JaxModelOp, which the code does not do. Workstream B1 closes the gap.

Stage 6 · An amortized retrieval (pipekit-train → Equinox)¶

TL;DR. We train a network that returns a mean and an uncertainty for the column enhancement for any measured spectrum, and we check its calibration.

Problem. For $\alpha \sim \mathcal{U}(0, 300)$ and the spectra $\vect{y} \in \mathbb{R}^5$ it produces, a network $g_\text{enc}(\vect{y}; \gvect{\nu}) \in \mathbb{R}^2$ outputs a Gaussian $q(\alpha \mid \vect{y}; \gvect{\nu}) = \mathcal{N}(\alpha; \mu_\alpha(\vect{y}; \gvect{\nu}), \sigma_\alpha^2(\vect{y}; \gvect{\nu}))$, with the mean $\mu_\alpha \in \mathbb{R}$ and the standard deviation $\sigma_\alpha \in \mathbb{R}_{>0}$. Its weights $\gvect{\nu} \in \mathbb{R}^{D_\nu}$ are variational parameters (they parameterise an approximate posterior), and they are trained on the negative log-likelihood (2),

A calibrated $q$ has $\mathbb{E}[r^2] = 1$ for the standardised residual $r = (\alpha - \mu_\alpha) / \sigma_\alpha \in \mathbb{R}$, and about 95 % of its $\pm 2 \sigma_\alpha$ intervals cover the truth. No unbiased estimator can beat the Cramér–Rao bound $\sigma_\mathrm{CR} = \sigma_\varepsilon / \Vert \vect{J}_h(\alpha) \Vert_2$, where $\vect{J}_h(\alpha) = \partial h / \partial \alpha \in \mathbb{R}^5$ is the Jacobian of the noise-free spectrum, with entries $[\vect{J}_h]_b = -k_b e^{-k_b \alpha}$. An RMSE averages squared errors, so its floor is the root-mean-square bound over the prior, $(\mathbb{E}_\alpha[\sigma_\mathrm{CR}^2])^{1/2} = 0.257$, whereas the mean bound, $\mathbb{E}_\alpha[\sigma_\mathrm{CR}] = 0.228$, is the comparator for a mean posterior width. On the 200 test draws below, the two are 0.253 and 0.227.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# Shapes: params = (α, ε₁…ε₅): (6,); y (5,); network 5 → 64 → 64 → 64 → 2
sim = pt.SimulationDataset(pc.CallableForward(beer_lambert, dt=1.0),
                           AlphaNoisePrior(), n_samples=20_000)        # today
train = Flip(sim)            # (params, y) → (y, α): a TrainingDataset subclass  user

class GaussianHead(eqx.Module):                                        # user
    mean: jax.Array
    log_var: jax.Array
    # log q(α | y) = −½ (log 2π + log σ_α² + (α − μ_α)² / σ_α²)
    def log_prob(self, a): ...

loop = pt.TrainingLoop(model_op=EquinoxModelOp(Net(key)),  # y → GaussianHead
                       dataset=train, val_dataset=train.with_split("val"),
                       loss=pt.NLL(), optimizer_config={"name": "adam",
                       "lr": 2e-3}, max_steps=8000, batch_size=256)   # today
q, _ = loop.run()                                          # 425 s on CPU
# 200 test draws: RMSE 0.303 (RMS Cramér–Rao floor 0.253); bias −0.15
# ±2σ_α coverage 99.5 %; E[r²] = 0.43: σ_α is about 1.5× too wide

The network is useful and honest in shape, as its $\sigma_\alpha$ grows from 0.33 for $\alpha < 100$ to 0.74 for $\alpha > 200$ just as the bound does. However, it is under-confident, since $\mathbb{E}[r^2] = 0.43$ means its intervals are about 1.5 times too wide. NLL only needs predicted.log_prob(target), which is the right seam, but there is no Gaussian head to plug into it, and the noise has to be drawn inside the prior so that it is seeded. Workstream D1 turns this into a recipe.

Stage 7 · The same retrieval with the Bayesian backends (pipekit-train → NumPyro, BlackJAX)¶

TL;DR. We solve one retrieval exactly with MCMC, as the reference that the amortized network should match.

Problem. For one measured spectrum $\vect{y} \in \mathbb{R}^5$ (with the true $\alpha = 120$), the posterior is (3),

which is nearly Gaussian here, with the Laplace standard deviation $\sigma_\varepsilon / \Vert \vect{J}_h(\alpha) \Vert_2 = 0.168$.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# Shapes: x = k (5, 1); y (5,); α scalar
def model(x, y=None):                                            # user
    a = numpyro.sample("alpha", dist.Uniform(0.0, 300.0))
    numpyro.sample("obs", dist.Normal(jnp.exp(-x[..., 0] * a), 1e-3), obs=y)

data = pt.IterableDataset(source=list(zip(k[:, None], y)),
                          content_hash="beer-lambert-y3")            # today
task = NumpyroTask(model, num_warmup=1000, num_samples=4000)         # today
for backend, t in [("numpyro-mcmc", task),
                   ("blackjax", BlackjaxTask(numpyro_task=task)),
                   ("numpyro-svi", task)]:
    op, art = pt.TrainingLoop(model_op=Dummy(), dataset=data, task=t,
                              backend=backend, max_steps=4000).run()  # today
# numpyro-mcmc: α = 120.143 ± 0.167, 0 divergences, 5.6 s
# blackjax:     α = 120.152 ± 0.168, 0 divergences, 10.4 s
# numpyro-svi:  α ≈ 119.96 ± 0.65 after 4,000 steps, 305 s

Both samplers agree with each other and with the Laplace width to the third decimal, so they are the oracle that we can score the amortized network of stage 6 against (its mean $\sigma_\alpha$ over the prior, 0.46, is close to three times this width). However, SVI, which minimises $\mathcal{L}_\text{ELBO} = -\mathrm{ELBO}$ (the negative evidence lower bound), is far from converged after 4,000 steps, and it is slow. This is because the adapter calls svi.update in a Python loop without jax.jit, at 76 ms per step on five data points. Workstream C4 jits the step. The Bayesian backends are not in the scoping report. They come from ADRs D13 and D14 (“inference is training too”), where NUTS is the oracle rung of the benchmark ladder.

Stage 8 · Train on stored states, and cache what you simulate (pipekit-train)¶

TL;DR. We read input and target windows from a Zarr archive of model states, and we cache the simulated pairs so that later runs skip the simulator.

Problem. For an archive of hourly states $\vect{x}_1, \ldots, \vect{x}_{50} \in \mathbb{R}^{D_x}$, a window dataset with an input stencil of the last three hours and a target one hour ahead yields the pairs $((\vect{x}_{t-2}, \vect{x}_{t-1}, \vect{x}_t), \vect{x}_{t+1})$ for every valid origin $t$. A cache keyed by the dataset’s content hash must return exactly what the source would.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# Shapes: archive (T = 50, D_x = 40); input (1, 3, D_x); target (1, 1, D_x)
stencils = {"input": TimeStencil("-2h", "0h", "1h", closed="both"),
            "target": TimeStencil("1h", "1h", "0h", closed="both")}
ds = pt.XarrayWindowDataset(source={"input": store, "target": store},
                            stencils=stencils, store_version="adv-v1")  # today
# → ImportError unless geopatcher.time is shimmed: geopatcher ships
#   geopatcher.temporal.stencils
len(ds), len(ds.with_split("val")), len(ds.with_split("test"))  # → 35, 1, 3
x, y = ds[0]                       # y = F x[0, −1] exactly            today

cached = pt.CachedDataset(sim_M, cache_dir="cache")                    # today
list(cached)                       # 400 items: 3.70 s, then 0.68 s
pt.CachedDataset(sim_2M, cache_dir="cache")[0]   # simulator now 2F …
# → returns the cached F pairs: the hash does not include the simulator

The window dataset is the training reader for the archives of Report 13. Its content_hash covers store_version, the stencils and the origins, and a linear last-step model trained on it reached an MSE of $3.5 \times 10^{-5}$. However, two defects stand out. It imports geopatcher.time, which the pinned geopatcher no longer has (xarray_window.py:35), and its error message wrongly says that geopatcher is not installed (A5). In addition, the 80/10/10 contiguous split with a guard band left one validation window out of 47 origins, which is the expected cost of guard bands on a short archive. The cache is correct when the source is, but SimulationDataset.content_hash hashes the forward model’s class name, dt and state signature and not its function, so swapping $\mat{F}$ for $2\mat{F}$ serves stale data (A3). It is also slow, because each __getitem__ reopens the zarr store, at 8.2 ms against 0.05 ms to recompute this simulator (A4).

The scoping report’s names, checked¶

Table 3 checks the stack names that the scoping report’s examples use.

Current-state audit¶

Of the twenty pieces we checked, ten ship and work, eight ship with defects, and two are missing (Table 4), and Figure 2 counts the page’s calls by marker. The audit read references/geostack/pipekit at ef3c17f on 13 October 2026.

There are also smaller findings from the design docs:

• The README says the package is “planning only” and that “the source directory currently contains no operators”, and the __init__ docstring still says “v0.0 ships the carrier-agnostic Protocols”.

• The ADRs are ahead of their status lines. D13 and D14 say “implementation planned”, blackjax_adapter.md says “Design only — no code yet”, and adapters.md heads BlackJAX “(planned)”, yet all three backends ship. decisions.md says “Eleven ADRs total” and lists 14.

• Hyperparameter search contradicts itself, as it is out of scope in vision.md, a non-goal in boundaries.md §12, and resolved and shipped in boundaries.md Q5.

• Checkpoints are described three ways: a Composite Orbax handler in adapters.md, PyGrainCheckpointHandler in D8, and the JSON side-car that the code actually writes (data_iter.json).

• IterableDataset and the sweep classes are undocumented in the API index, and the layout in architecture.md omits xarray_window.py, _xreader.py and the Bayesian adapters.

• Seed ranges overlap across datasets. The splits within one SimulationDataset are disjoint, but two datasets with seed=0 and seed=1 share all but one sample seed.

Workstream A · Datasets¶

Decision. Each of the three training shapes gets a dataset that is correct, hashed by everything that determines its contents, and fast enough to feed Grain (Table 5).

A1 · CatalogDataset¶

1. Read from a catalog. It iterates the rows that a GeoSlice query returns, or the patches that a geopatcher Sampler (such as RegularStride) cuts from them, through geocatalog’s CatalogDomain bridge.

2. Reuse the inference preprocess. preprocess is the same operator that the inference pipeline runs, so training and serving cannot drift apart.

3. Hash the recipe. content_hash covers the catalog’s own hash (a CatalogBundle or the GeoParquet file’s checksum), the query, the sampler config, the preprocess config, the split and the seed.

4. Accept states too. With the state profile of Report 13, the same class reads a state catalog through load_states, so there is no separate StateCatalog type to support.

1
2
3
4
5
6
7
8
# Shapes: chips (N_batch = 4, 256, 256) → labels (2, 256, 256)
dataset = pt.CatalogDataset(                                     # proposed (A1)
    catalog=gc.open_catalog("s3://imeo/labeled-scenes.parquet"),        # today
    preprocess=gz.radiometry.ToFloat32() | gz.radiometry.PercentileClip(
        p_min=2, p_max=98),                                             # today
    target_op=load_label,                                    # arguments sketched
    sampler=gp.RegularStride((256, 256)),                    # arguments sketched
)

A2 · A correct simulation path¶

SimulationDataset calls cycle(params, state0) with a state that the user supplies (cycle_state=). It checks once, by signature, that the prior accepts seed, rather than catching TypeError on every call, and it documents that the observation noise is drawn inside the prior so that it is seeded.

A3 · Hash the simulator¶

content_hash includes the forward model’s get_config() and, for a CallableForward, a user-supplied version string, since a Python function cannot be hashed reliably. Without either, it warns that the cache cannot tell two simulators apart.

A4 · A cache that keeps up¶

CachedDataset opens the zarr group once per process and reads in chunks. Parquet lands for tabular pairs, whereas tfrecord stays deferred unless a Lightning or Keras user needs it.

A5 · The window dataset’s import¶

We import TimeStencil from geopatcher.temporal.stencils (or from the geopatcher root, if it re-exports it), and we make the error message say which module is missing.

Workstream B · The trained model’s handoff¶

Decision. Every trained operator can be stored with its weights and reloaded to the same outputs, and the artifact always names it (Table 6). B1 is absorbed by Report 16 B1, which makes JaxModelOp the one Equinox wrapper and aliases EquinoxModelOp to it.

1
2
3
4
5
6
7
8
9
# Shapes: W (D_x, D_x)
# Today:
trained, art = loop.run()                       # EquinoxModelOp      today
j = JaxModelOp(trained.module)                  # manual re-wrap      user
h = reg.store(j, weights=j.serialize_weights())                       # today

# After:
trained, art = loop.run()                       # JaxModelOp     proposed (B1)
loaded = reg.load(art.trained_model_hash)       # with weights   proposed (B1, B2)

Workstream C · Loop components¶

Decision. Checkpoints, sweeps, validation and the SVI backend do what their names promise (Table 7).

Workstream D · Amortized inference¶

Decision. pipekit-train documents one recipe for amortized inference and ships the smallest piece that every user rewrote, whereas the density estimators stay in the model libraries (Table 8).

Workstream E · The Lightning and Keras backends¶

Decision. We build the two backends that ADR D5 already scheduled, on the same run(loop) -> (op, backend_info) seam as Equinox (Table 9).

Testing gates¶

Each workstream only merges when its gates pass, and the existing gates are checks that the code passes today (Table 10).

Sequencing and risks¶

The fixes come first, because the running example hit them. The missing features follow, and the two new backends come last, as ADR D5 already orders them (Figure 3).

• CatalogDataset waits on packaging. The [catalog] extra is held back until geocatalog is on PyPI. Fallback: we ship CatalogDataset without an extra and import geocatalog lazily with an install hint, as XarrayWindowDataset does for geopatcher.

• A simulator hash cannot see inside a function. Hashing the code of a CallableForward is fragile across Python versions. Fallback: a required version string for callables, with a warning when it is absent.

• Changing the returned operator type breaks callers. Scripts that use trained.module on an EquinoxModelOp must keep working. Fallback: we keep EquinoxModelOp and give it serialize_weights and from_registry instead of returning JaxModelOp.

• Two more backends to maintain. Lightning and Keras double the test matrix. Fallback: we mark both as experimental and gate them only on the emulator gate E1/E2, as the scoping report itself suggested.

• Under-confident posteriors look safe. The intervals of stage 6 cover 99.5 %, so they would pass a naive coverage check. Fallback: the recipe reports $\mathbb{E}[r^2]$ and compares against an MCMC oracle on a few scenes, as stage 7 does.

Open questions¶

1. Where do distribution heads and density estimators live? The design docs name pyrox for ConditionalNormalizingFlow, but it is not pinned here. Is a GaussianHead a pyrox object, or a small pipekit_train.recipes helper?

2. Should a trained model’s prior travel with it? The emulator of stage 3 is only accurate inside its training prior. Should the artifact record the prior’s hash, and should NeuralForward warn on inputs outside it?

3. Seeds across datasets. Two SimulationDatasets with seeds 0 and 1 share all but one sample seed. Should the seeds be hashed into the sample seed rather than added?

4. One validation path or two? The Equinox loop evaluates with the task’s loss, whereas ValidationStep evaluates with arbitrary metrics. C3 makes the loop use ValidationStep, so is the task-loss path still needed?

5. Sweeps in parallel. HyperSweep runs the trials one after another, and parallel trials and Optuna or Ray Tune samplers are deferred. Is pipekit_ops.parallel.ProcessMap enough for the first parallel version?

6. Online training. OnlineTrainer is deferred to v0.3 (boundaries Q6), and the scoping report ruled online training out. Is it still wanted?

Upcoming work¶

Table 11 lists the follow-ups.

References¶

Sources¶

• The scoping report, “Report 11 — pipekit-train: training pipelines for emulators and amortized inference” (supplied by the author).

• pipekit-train 0.0.2 in references/geostack/pipekit/packages/pipekit-train, with its source, tests, README, CHANGELOG and docs/design/ (vision, architecture, boundaries, decisions D1 to D14, API pages and examples).

• The running example’s scripts and outputs, described in Section ??.