pipekit-train

Summary¶

Decision. pipekit-train stays a thin layer over existing training tools, as both the scoping report and ADR D1 say, and its 0.0.2 design stands: 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 supports only zarr (A4). XarrayWindowDataset imports a geopatcher module that no longer exists (A5).

• B. The trained model’s handoff. The Equinox backend returns an EquinoxModelOp that cannot be stored with its weights or reloaded from LocalModelRegistry; wrapping it in pipekit-jax’s JaxModelOp works (B1). 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 validation (C2); ValidationStep is never used by the loop (C3); 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. Ship a recipe and a Gaussian head, and decide where conditional density estimators live (D1).

• E. The other two backends. 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: supervised training from a catalog, emulator training from a pipekit-cycle forward model, and amortized inference. Each 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 data, a model and a training tool; the tools own the training math and the model libraries own the models (Figure 1).

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

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 Report 10’s two problems, so that a trained model can be checked against the exact operator it replaces. In the six-step cycle, stage 2 is step 3 of a forecast (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 Report 10’s upwind advection on a periodic 40-cell transect, $x_{t+1} = M x_t$ with $M = 0.6 I + 0.4 S$. An emulator is a learned $\hat{M}$; the target is $\hat{M} = M$.

The retrieval uses the five-band sensor of Report 17 unchanged ($y_b = e^{-k_b \alpha} + \varepsilon_b$, $k = (0.002, \dots, 0.010)$, $\sigma = 10^{-3}$, $\alpha \in [0, 300]$). An amortized retrieval learns $p(\alpha \mid y)$ once, for every $y$; a per-scene retrieval (Report 10’s Gauss–Newton, or MCMC) solves for one $y$ at a time.

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

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

TL;DR. Turn Report 10’s transport model into 2,000 input–output pairs with fixed seeds and train, validation and test splits.

Problem. A SimulationDataset draws $x^{(i)}$ from a prior $\pi$ with seed $s_0 + o + i$, where $o$ is the split’s offset, and returns $(x^{(i)}, M x^{(i)})$. 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: N = 40; 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)    # (N,) → (N,)    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. 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, so every sample comes out identical. Workstream A2 fixes both.

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

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

Problem. With pairs $(x^{(i)}, M x^{(i)})$ and a linear model $\hat{M} = W$, the mean-squared error (1)

is zero at $W = M$, and the minimiser is unique only when the inputs span $\mathbb{R}^N$, that is when $X = [x^{(1)}, \ldots, x^{(n)}]$ has full rank.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# Shapes: W (N, N); 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 − M| = 3.8e-7
# smooth prior: stops at step 2200; val MSE 3.4e-5;   max |W − M| = 0.343

Both runs fit their data (Table 2). The smooth prior’s inputs nearly lie in a low-dimensional subspace, so $W$ is fixed only on that subspace and is arbitrary off it: a low training loss is not a correct emulator. This is a property of the data, 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. Replace the forward model of a six-hour forecast with the emulator, and test it on states the prior did and did not cover.

Problem. The emulated forecast is $\hat{x}_6 = W^6 x_0$; the error $\lVert W^6 x_0 - M^6 x_0 \rVert_\infty$ is small only when $x_0$ lies where $W$ was fixed.

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

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

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

TL;DR. 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: 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. Checkpoint has two defects. It writes to the loop’s checkpoint_dir and ignores its own save_dir (the adapter reads only every_n_steps and keep_last). And the final step is saved only 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. Put the trained emulator in the registry and load it back with its weights.

Problem. A registered model must reload to the same outputs: after store and load, $\max \lvert W_{\text{loaded}} - M \rvert$ should equal the trained $3.8 \times 10^{-7}$.

1
2
3
4
5
6
7
8
9
10
11
12
# Shapes: W (N, N)
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 − M| = 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 the Equinox adapter returns a JaxModelOp; the code does not. Workstream B1 closes the gap.

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

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

Problem. For $\alpha \sim \mathcal{U}(0, 300)$ and the spectra $y$ it produces, a network $g_\phi$ outputs a Gaussian $q_\phi(\alpha \mid y) = \mathcal{N}(\mu_\phi(y), \sigma_\phi^2(y))$ and is trained on the negative log-likelihood (2),

A calibrated $q_\phi$ has $\mathbb{E}[z^2] = 1$ for $z = (\alpha - \mu_\phi) / \sigma_\phi$ and covers the truth with about 95 % of its $\pm 2 \sigma$ intervals. No unbiased estimator can beat the Cramér–Rao bound $\sigma_{\mathrm{CR}} = \sigma / \lVert J(\alpha) \rVert$ with $J_b = -k_b e^{-k_b \alpha}$. An RMSE averages squared errors, so its floor is the root-mean-square bound over the prior, $\left(\mathbb{E}_\alpha\left[\sigma_{\mathrm{CR}}^2\right]\right)^{1/2} = 0.257$; 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: θ = (α, ε₁…ε₅): (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)            # (θ, 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[z²] = 0.43: σ̂ is about 1.5× too wide

The network is useful and honest in shape (its $\hat{\sigma}$ grows from 0.33 for $\alpha < 100$ to 0.74 for $\alpha > 200$, as the bound does) but under-confident: $\mathbb{E}[z^2] = 0.43$ means its intervals are about 1.5 times too wide. NLL needs only 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. Solve one retrieval exactly with MCMC, as the reference the amortized network should match.

Problem. For one measured $y$ (true $\alpha = 120$), the posterior is (3),

which is nearly Gaussian here, with the Laplace standard deviation $\sigma / \lVert J(\alpha) \rVert = 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: they are the oracle the amortized network of stage 6 can be scored against (its mean $\hat{\sigma}$ over the prior, 0.46, is close to three times this width). SVI is far from converged after 4,000 steps, and slow: 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 benchmark ladder’s oracle rung.

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

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

Problem. For an archive $x_1, \ldots, x_{50}$, a window dataset with an input stencil of the last three hours and a target one hour ahead yields $((x_{t-2}, x_{t-1}, x_t), 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, N = 40); input (1, 3, N); target (1, 1, N)
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 = M · 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 2M …
# → returns the cached M pairs: the hash does not include the simulator

The window dataset is the training reader for Report 13’s archives: its content_hash covers store_version, the stencils and the origins, and a linear last-step model trained on it reached MSE $3.5 \times 10^{-5}$. 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 geopatcher is not installed (A5). And 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, not its function, so swapping $M$ for $2M$ serves stale data (A3). It is also slow: each __getitem__ reopens the zarr store, 8.2 ms against 0.05 ms to recompute this simulator (A4).

The scoping report’s names, checked¶

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

Current-state audit¶

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

Smaller findings from the design docs:

• The README says the package is “planning only” and that “the source directory currently contains no operators”; 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)”; all three backends ship. decisions.md says “Eleven ADRs total” and lists 14.

• Hyperparameter search contradicts itself: out of scope in vision.md, a non-goal in boundaries.md §12, 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 the code actually writes (data_iter.json).

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

• Seed ranges overlap across datasets. 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. The three training shapes each get 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. Iterate the rows a GeoSlice query returns, or the patches a geopatcher Sampler (such as RegularStride) cuts from them, through geocatalog’s CatalogDomain bridge.

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

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 Report 13’s state profile, the same class reads a state catalog through load_states; there is no separate StateCatalog type to support.

1
2
3
4
5
6
7
8
# Shapes: chips (B = 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 the user supplies (cycle_state=), checks that the prior accepts seed by signature once rather than by catching TypeError on every call, and documents that observation noise is drawn inside the prior so 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, and tfrecord stays deferred unless a Lightning or Keras user needs it.

A5 · The window dataset’s import¶

Import TimeStencil from geopatcher.temporal.stencils (or the geopatcher root, if it re-exports it), and 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 (N, N)
# 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 every user rewrote; density estimators stay in the model libraries (Table 8).

Workstream E · The Lightning and Keras backends¶

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

Testing gates¶

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

Sequencing and risks¶

Fixes come first, because the running example hit them; the missing features follow, and the two new backends 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: ship CatalogDataset without an extra, importing geocatalog lazily with an install hint, as XarrayWindowDataset does for geopatcher.

• A simulator hash cannot see inside a function. Hashing a CallableForward’s code 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: 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: mark both experimental and gate them only on the emulator gate E1/E2, as the scoping report itself suggested.

• Under-confident posteriors look safe. Stage 6’s intervals cover 99.5 % and would pass a naive coverage check. Fallback: the recipe reports $\mathbb{E}[z^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? pyrox is named in the design docs 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? Stage 3’s emulator is accurate only 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 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; ValidationStep evaluates with arbitrary metrics. C3 makes the loop use ValidationStep; is the task-loss path still needed?

5. Sweeps in parallel. HyperSweep runs trials one after another; parallel trials and Optuna or Ray Tune samplers are deferred. Is pipekit.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: its source, tests, README, CHANGELOG and docs/design/ (vision, architecture, boundaries, decisions D1–D14, API pages and examples).

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