Skip to content

pipekit-train Design Doc

A thin orchestration layer over existing training tools, exposed as pipekit.Operators. Trained models drop into inference pipelines without rewrite.

Structure

docs/design/
├── README.md              # This file — index and reading order.
├── vision.md              # Motivation, scope, three training shapes.
├── architecture.md        # The layered design + integration with
│                          # pipekit, pipekit-cycle, pipekit-experiment.
├── boundaries.md          # Non-goals, open questions, future work.
├── decisions.md           # ADRs D1–D14 (carrier-agnostic + Equinox +
│                          # sharding + NumPyro + BlackJAX).
├── numpyro_adapter.md     # Full design — the NumPyro backends numpyro-svi
│                          # + numpyro-mcmc. Companion to api/adapters.md.
├── blackjax_adapter.md    # Full design — the BlackJAX sampler backend
│                          # (NUTS / SG-MCMC). Sibling to numpyro_adapter.md.
├── api/
│   ├── README.md          # Core abstractions (TrainingDataset,
│   │                      # Loss, TrainingLoop, Callback, MetricWriter).
│   ├── datasets.md        # Layer 0 — TrainingDataset, CatalogDataset,
│   │                      # SimulationDataset, CachedDataset.
│   ├── primitives.md      # Layer 1 — Loss protocol, common Loss ops,
│   │                      # train_step (Equinox), save/restore (Orbax).
│   ├── components.md      # Layer 2 — Callback, MetricWriter,
│   │                      # ValidationStep, EarlyStopping, Checkpoint,
│   │                      # LogToExperiment.
│   ├── models.md          # Layer 3 — TrainingLoop (StatefulOperator),
│   │                      # the outer loop, the TrainingArtifact handshake.
│   └── adapters.md        # Layer 4 — Equinox (v0.1 reference),
│                          # Lightning (v0.2 scaffold), Keras (v0.3 scaffold).
└── examples/
    ├── README.md          # Index and reading order.
    ├── direct.md          # Cloud-mask classifier (CatalogDataset).
    ├── emulator.md        # Chemistry surrogate (SimulationDataset).
    ├── amortized.md       # Neural posterior for plume source (SBI).
    └── integration.md     # With pipekit-cycle, pipekit-experiment,
                           # geocatalog, geopatcher, somax, vardax, pyrox.

Reading Order

  1. vision.md — what gap this fills, what it isn't.
  2. architecture.md — the layered design and where each piece sits in the stack.
  3. decisions.md — the ADRs that lock the shape in.
  4. api/README.md → drill into datasets → primitives → components → models → adapters.
  5. examples/ — see the three training shapes in code.
  6. boundaries.md — what's deferred and why.

Sources

The design synthesises two upstream sources:

  • The master plan: research_journal_v2/notes/geotoolz/master_plan/toolz_9_pipekit_train.md (Report 11). Vision, three training shapes, package scope, reproducibility story.
  • The eqx_trainer design doc in jej_vc_snippets/design_docs/eqx_trainer/. Concrete JAX-native stack (Equinox + Optax + Grain + Orbax) that becomes the v0.1 reference adapter inside pipekit_train.adapters.equinox.

Where the two diverge, this doc resolves the tension explicitly — see decisions.md for the ADRs.