Prune the uv cache at the end of every CI job

The uv cache lives on a persistent volume shared by every job on the runner (/srv/act-runner-cache/uv), so nothing trimmed it and it grew without bound. `uv cache prune --ci` drops the entries that are not worth keeping between runs (pre-built wheels for local sources) while leaving the downloaded-wheel cache that makes `uv sync` fast. In bump-version the step carries the same is_merge guard as the rest of that job, since uv is only set up on a merge push. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Seed each epoch's RNG from (seed, epoch) (gitea #83 )
2026-08-26 13:08:26 +02:00 · 2026-08-26 12:36:59 +02:00 · 2026-08-26 12:36:59 +02:00 · 2026-08-26 12:36:59 +02:00 · 2026-08-26 12:36:59 +02:00 · 2026-08-26 08:15:52 +00:00
127 changed files with 19095 additions and 9072 deletions
@@ -0,0 +1,17 @@
+[tool.bumpversion]
+current_version = "0.3.10"
+parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
+serialize = ["{major}.{minor}.{patch}"]
+search = "{current_version}"
+replace = "{new_version}"
+regex = false
+allow_dirty = false
+commit = true
+tag = false
+message = "chore: bump version {current_version} -> {new_version} [skip ci]"
+pre_commit_hooks = ["uv lock", "git add uv.lock"]
+
+[[tool.bumpversion.files]]
+filename = "pyproject.toml"
+search = "version = \"{current_version}\""
+replace = "version = \"{new_version}\""
@@ -28,6 +28,9 @@ jobs:
          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
      - run: uv sync --extra cpu --extra dev
      - run: uv run ruff check .
+      # The uv cache is a persistent volume shared by every job on this
+      # runner, so each job trims what it no longer needs before exiting.
+      - run: uv cache prune --ci

  ruff-format:
    name: Format (ruff format)
@@ -46,6 +49,7 @@ jobs:
          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
      - run: uv sync --extra cpu --extra dev
      - run: uv run ruff format --check .
+      - run: uv cache prune --ci

  type-check:
    name: Type check (ty)
@@ -64,6 +68,7 @@ jobs:
          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
      - run: uv sync --extra cpu --extra dev
      - run: uv run ty check .
+      - run: uv cache prune --ci

  test:
    name: Tests
@@ -82,7 +87,102 @@ jobs:
          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
      - run: uv sync --extra cpu --extra dev
-      - run: uv run pytest
+      - run: uv run pytest --cov --cov-report=term-missing --cov-report=xml
+      - uses: actions/upload-artifact@v3
+        with:
+          name: coverage-report
+          path: coverage.xml
+      - run: uv cache prune --ci
+
+  bump-version:
+    name: Bump version, tag, and update changelog on merge to master
+    needs: [ruff-check, ruff-format, type-check, test]
+    if: github.ref == 'refs/heads/master' && github.event_name == 'push'
+    runs-on: ubuntu-latest
+    container:
+      image: docker.gitea.com/runner-images:ubuntu-latest
+      volumes:
+        - /srv/act-runner-cache/uv:/uv-cache
+    steps:
+      # CI_TOKEN needs write:repository scope (not just read) — this job
+      # pushes commits and tags to master, unlike ruff-check/ruff-format/
+      # type-check/test above, which only need to check out the repo.
+      - uses: actions/checkout@v4
+        with:
+          token: ${{ secrets.CI_TOKEN }}
+          fetch-depth: 0
+      - name: Check whether this push is a merge commit
+        id: merge_check
+        run: |
+          PARENTS=$(git rev-parse HEAD^@ | wc -l)
+          echo "HEAD has $PARENTS parent(s)"
+          if [ "$PARENTS" -ge 2 ]; then
+            echo "is_merge=true" >> "$GITHUB_OUTPUT"
+          else
+            echo "is_merge=false" >> "$GITHUB_OUTPUT"
+          fi
+      - uses: astral-sh/setup-uv@v5
+        if: steps.merge_check.outputs.is_merge == 'true'
+        with:
+          enable-cache: false
+      - run: |
+          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
+          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
+        if: steps.merge_check.outputs.is_merge == 'true'
+      - run: uv sync --extra cpu --extra dev
+        if: steps.merge_check.outputs.is_merge == 'true'
+      - name: Configure git identity
+        if: steps.merge_check.outputs.is_merge == 'true'
+        run: |
+          git config user.name "gitea-actions"
+          git config user.email "actions@git.larsbogner.de"
+      - name: Bump patch version if this merge didn't already bump it
+        if: steps.merge_check.outputs.is_merge == 'true'
+        run: |
+          OLD_VERSION=$(git show "${{ github.event.before }}:pyproject.toml" 2>/dev/null | grep -m1 '^version = ' | sed -E 's/version = "(.*)"/\1/')
+          CURRENT_VERSION=$(uv version --short)
+          if [ -z "$OLD_VERSION" ]; then
+            echo "Could not read pyproject.toml at github.event.before; falling back to HEAD^1"
+            OLD_VERSION=$(git show "HEAD^1:pyproject.toml" | grep -m1 '^version = ' | sed -E 's/version = "(.*)"/\1/')
+          fi
+          if [ "$OLD_VERSION" = "$CURRENT_VERSION" ]; then
+            echo "Version unchanged by this merge ($CURRENT_VERSION); bumping patch"
+            uv run bump-my-version bump patch --current-version "$CURRENT_VERSION"
+          else
+            echo "Branch already bumped the version ($OLD_VERSION -> $CURRENT_VERSION); skipping auto-bump"
+          fi
+      - name: Update changelog for the current version if not already tagged
+        if: steps.merge_check.outputs.is_merge == 'true'
+        run: |
+          VERSION=$(uv version --short)
+          TAG="v$VERSION"
+          if git rev-parse "$TAG" >/dev/null 2>&1; then
+            echo "Tag $TAG already exists; skipping changelog update"
+          else
+            uv run git-cliff --tag "$TAG" --unreleased --prepend CHANGELOG.md
+            git add CHANGELOG.md
+            if ! git diff --cached --quiet -- CHANGELOG.md; then
+              git commit -m "chore: update changelog for $TAG [skip ci]"
+            else
+              git restore --staged CHANGELOG.md
+            fi
+          fi
+      - name: Push commits and tag the current version
+        if: steps.merge_check.outputs.is_merge == 'true'
+        run: |
+          git push origin HEAD:master
+          VERSION=$(uv version --short)
+          TAG="v$VERSION"
+          if git rev-parse "$TAG" >/dev/null 2>&1; then
+            echo "Tag $TAG already exists"
+          else
+            git tag -a "$TAG" -m "$TAG"
+            git push origin "refs/tags/$TAG"
+          fi
+      # Same guard as every other step here: on a non-merge push uv was never
+      # set up, so there is nothing to prune.
+      - run: uv cache prune --ci
+        if: steps.merge_check.outputs.is_merge == 'true'

  sync-version-on-tag:
    name: Sync project version with tag
@@ -111,3 +211,4 @@ jobs:
          else
            echo "Tag version matches project version ($CURRENT_VERSION)"
          fi
+      - run: uv cache prune --ci
@@ -20,3 +20,8 @@ checkpoints/

 # giant analyze run directories (shared.json, reduced/, plots/, condor logs)
 /analysis_runs/
+
+# Coverage artifacts
+.coverage
+coverage.xml
+htmlcov/
@@ -0,0 +1,551 @@
+# Changelog
+
+## [0.3.10] - 2026-08-26
+
+### Changed
+
+- Backfill CHANGELOG.md for v0.2.0-v0.3.2
+
+- Docs: bring README and CLAUDE.md in line with v0.3.9
+
+## [0.3.9] - 2026-08-24
+
+### Added
+
+- Add multi-rollout support to giant analyze [gitea #77](https://git.larsbogner.de/lars/giant/issues/77)
+
+
+### Changed
+
+- Escape LaTeX-special characters in plot titles/xlabels [gitea #81](https://git.larsbogner.de/lars/giant/issues/81)
+
+## [0.3.8] - 2026-08-24
+
+### Added
+
+- Add giant analyze metrics plots for training progress [gitea #75](https://git.larsbogner.de/lars/giant/issues/75)
+
+
+### Fixed
+
+- Fix LaTeX-unavailable skip check in analyze metrics smoke test
+
+## [0.3.7] - 2026-08-24
+
+### Added
+
+- Add rollout-quality distance, confusion, containment and router plots [gitea #76](https://git.larsbogner.de/lars/giant/issues/76)
+
+## [0.3.6] - 2026-08-24
+
+### Changed
+
+- Give CriticModel a registry-built trunk and StageModel base [gitea #57](https://git.larsbogner.de/lars/giant/issues/57)
+
+## [0.3.5] - 2026-08-24
+
+### Added
+
+- Add "none" variants for router, history, and trunk [gitea #45](https://git.larsbogner.de/lars/giant/issues/45)
+
+## [0.3.4] - 2026-08-23
+
+### Added
+
+- Add giant model summary command [gitea #46](https://git.larsbogner.de/lars/giant/issues/46)
+
+- Add per-stage init_from/freeze [gitea #42](https://git.larsbogner.de/lars/giant/issues/42)
+
+- Add bf16 autocast to the training loop [gitea #47](https://git.larsbogner.de/lars/giant/issues/47)
+
+- Add class-balanced secondary particle-type loss [gitea #44](https://git.larsbogner.de/lars/giant/issues/44)
+
+
+### Changed
+
+- Implement stage2_model.stage1_context = "sampled" [gitea #41](https://git.larsbogner.de/lars/giant/issues/41)
+
+- Bump patch version to 0.3.3
+
+- Offset event_id across multi-shard reference reads in giant analyze [gitea #22](https://git.larsbogner.de/lars/giant/issues/22)
+
+- Auto-bump patch version, tag, and update changelog on merge to master [gitea #50](https://git.larsbogner.de/lars/giant/issues/50)
+
+- Document CI_TOKEN's write:repository scope requirement [gitea #50](https://git.larsbogner.de/lars/giant/issues/50)
+
+## [0.3.2] - 2026-08-17
+
+### Added
+
+- Add configs/baseline.toml as the kept reference model
+
+
+### Fixed
+
+- Clamp analysis histogram bins before the i32 cast, not after [gitea #61](https://git.larsbogner.de/lars/giant/issues/61)
+
+- Clip raw predicted log_mass in decode_secondaries [gitea #54](https://git.larsbogner.de/lars/giant/issues/54)
+
+
+### Changed
+
+- Let dwarf warm-cache take --config so it can't under-warm a config's cache keys [gitea #59](https://git.larsbogner.de/lars/giant/issues/59)
+
+- Implement n_sec.mode = "stop_token" for the AR secondary decoder [gitea #40](https://git.larsbogner.de/lars/giant/issues/40)
+
+- Bump patch version to 0.3.2
+
+## [0.3.1] - 2026-08-14
+
+### Added
+
+- Add an Objective registry for the flow/ddpm/wgan generator choice [gitea #32](https://git.larsbogner.de/lars/giant/issues/32)
+
+
+### Changed
+
+- Make trunk architecture selectable via a registry [gitea #33](https://git.larsbogner.de/lars/giant/issues/33)
+
+- Make ResBlock's conditioning-injection mechanism selectable [gitea #34](https://git.larsbogner.de/lars/giant/issues/34)
+
+- Make HistoryEncoder a pluggable registry, like Router/Objective [gitea #35](https://git.larsbogner.de/lars/giant/issues/35)
+
+- Deduplicate n_sec_head/type_head MLPs into build_mlp_head [gitea #36](https://git.larsbogner.de/lars/giant/issues/36)
+
+- Give the cond_cat/cond_cont column layout one owner [gitea #37](https://git.larsbogner.de/lars/giant/issues/37)
+
+- Give Stage1Model/Stage2OneShot/Stage2Autoregressive a shared StageModel base [gitea #39](https://git.larsbogner.de/lars/giant/issues/39)
+
+- Pass ConditioningAxisConfig/ParticleTypeConfig themselves instead of raw dicts [gitea #38](https://git.larsbogner.de/lars/giant/issues/38)
+
+- Bump patch version to 0.3.1
+
+## [0.3.0] - 2026-08-13
+
+### Added
+
+- Add v0.3.0 design doc: Stage-2 autoregressive redesign
+
+- Add pytest-cov to dev deps and run coverage in CI
+
+- Add coverage for router-center seeding, geometry batch reader, material topN cache, and setup-cache corruption paths
+
+- Add render.py coverage: figure params, router diagnostics plots, gallery/condor glue
+
+- Add unknown-key validation to config.toml merge (issues.md Issue 2)
+
+- Add consumed-keys audit test (issues.md Issue 5)
+
+
+### Fixed
+
+- Fix test_render_all_run_gallery_invokes_subprocess clobbering LaTeX's own subprocess.run
+
+
+### Removed
+
+- Remove issues.md
+
+
+### Changed
+
+- Refine v0.3.0 design: defaults, deferred scope, open questions
+
+- Document the differentiability position and its validation obligation
+
+- V0.3.0 step 1: new nested config schema, v0.2 migration shim
+
+- V0.3.0 step 2: network.py refactor to composable stage models
+
+- V0.3.0 step 3: per-stage train.py trainers + pipeline.py/cli.py rewrite
+
+- V0.3.0 step 4: type map + particle_type.target = "onehot"/"embedding"
+
+- V0.3.0 step 5: Stage2Autoregressive (history=markov) + §11.4 grad instrumentation
+
+- V0.3.0 step 6: sample.py/rollout.py AR generation + class->PDG decode
+
+- V0.3.0 step 7: AttentionHistory (KV-cached) + scheduled/never teacher forcing
+
+- V0.3.0 post-implementation audit: resolve all 9 tracked discrepancies
+
+- Refactor train.py into giant/training/ around a metrics collector
+
+- Silence the fork-safety warning from num_workers>0 pipeline tests
+
+- Deduplicate giant/training/trainers.py shared per-stage logic
+
+- Rewrite README for v0.3.0 architecture, quick start, and data columns
+
+- Bump version to 0.3.0
+
+- Delete docs/v0.3.0-design.md and strip all references to it
+
+- Apply ruff format
+
+- Downgrade coverage-report upload to actions/upload-artifact@v3
+
+- Bump ruff line-length to 120 and reformat
+
+- Make config dataclasses the single source of truth for DEFAULT_CONFIG
+
+- Extract giant train/new-run's CLI override mapping into a table-driven function (issues.md Issues 3 & 4)
+
+- Mark issues.md Issues 3 & 4 as fixed
+
+- Extract predict/rollout's duplicated inference bootstrap into giant.checkpoint_io (issues.md Issue 5)
+
+- Mark issues.md Issue 5 as fixed
+
+- Unify the two v0.2->v0.3 migration surfaces (issues.md Issue 6)
+
+- Type the data/model/training batch contracts with NamedTuples (issues.md Issue 7)
+
+- Split giant/model/network.py into giant/model/ (issues.md Issue 8)
+
+- Move scripts/ to giant/tools/ (issues.md Issue 9)
+
+- Reject stage2_model.stage1_context = 'sampled' as unimplemented (issues.md Issue 1)
+
+- Honour wgan.critic_hidden_dim/critic_n_res_blocks in build_critics [gitea #28](https://git.larsbogner.de/lars/giant/issues/28)
+
+- Validate stage2_model.autoregressive.order in validate_config [gitea #30](https://git.larsbogner.de/lars/giant/issues/30)
+
+- Decouple secondary-species vocabulary from conditioning.particle.emb_dim [gitea #29](https://git.larsbogner.de/lars/giant/issues/29)
+
+- Skip router auxiliary loss compute when their lambda is 0 [gitea #31](https://git.larsbogner.de/lars/giant/issues/31)
+
+## [0.2.0] - 2026-08-04
+
+### Added
+
+- Add CLAUDE.md with architecture overview and dev commands
+
+- Add streaming data pipeline and giant CLI entry point
+
+- Add giant predict command
+
+- Add ROOT-to-parquet conversion script with convert dependency group
+
+- Add post_pos as a model target via travel_dir decomposition
+
+- Add --coord local mode to predict for raw-space prediction debugging
+
+- Add KL divergence to marginal validation and hook it into the training loop
+
+- Add graceful shutdown on SIGINT/SIGTERM
+
+- Add configurable dropout to ResBlocks
+
+- Add giant.analysis module for notebook-based model quality diagnostics
+
+- Add lazy polars I/O and duplicate KL/constraint checks for giant.analysis
+
+- Add ruff and ty as dev dependencies, fix lint/type findings
+
+- Add linear warmup before cosine LR decay
+
+- Add --batch-size auto to estimate batch size from free GPU memory
+
+- Add hyperparameter scan
+
+- Add --batch-size auto to predict, matching train
+
+- Add tqdm progress bar to predict
+
+- Add KL bar plots and sample_frac to load_predicted_local; ignore root parquet scratch files
+
+- Add event-level shower observables to giant.analysis
+
+- Add total length traveled per event to event observables
+
+- Add pdg energy/length contribution pie plots
+
+- Add export script for Tier 4 event-level/pdg-share plots
+
+- Add mean/median deposited energy and step length plots per event
+
+- Add export script for ETP group-update presentation plots
+
+- Add photon edep export scripts and per-step presentation plots
+
+- Add tooling for a versioned geant_steps dataset layout
+
+- Add --copy mode to migrate_geant_steps.py
+
+- Add update-manifest and create-manifest subcommands to bump_dataset_version
+
+- Add --to flag for bump-gen/bump-schema and --gen flag for update-manifest
+
+- Add disk usage summary to dwarf status
+
+- Add file counts and reference tracking to dwarf status
+
+- Add --comment option to predict, recorded in YAML sidecar
+
+- Add energy-conservation PoC ODE-step comparison scripts
+
+- Add autoregressive shower rollout driver
+
+- Add fast slab lookup for the GeometryOracle, replacing knn as the default
+
+- Add load_rollout_vs_truth to compare rollouts against held-out truth data
+
+- Add mixture-of-experts routing prototype for Stage 1 and Stage 2
+
+- Add ProcessRouter for physics-process-based expert gating
+
+- Add PdgRouter for particle-type-based expert gating
+
+- Add ComposedRouter for multi-axis MoE gating
+
+- Add EMA weights, weight decay, step-based LR schedule, and grad-norm logging to training
+
+- Add WGAN-GP mode as a throwaway fast-eval experiment
+
+- Add router gating diagnostic for MoE checkpoints
+
+- Add Gitea Actions CI pipeline
+
+- Add configs for router energy (embedding/physical) and WGAN baseline runs
+
+- Add opt-in Weights & Biases logging for the training loop
+
+- Add test coverage for resolve_expert_dims
+
+- Add regression coverage for vocab/process index-map builders
+
+- Add dwarf warm-cache to precompute the setup-stage sidecar
+
+- Add giant new-run to scaffold a config.toml + run dir ahead of training
+
+- Add learnable per-expert width and shared temperature to EnergyRouter
+
+- Add opt-in straight-through Gumbel-softmax combine weights to MoE router
+
+- Add gumbel router configs sweeping learnable-knob combinations
+
+- Add gumbel/learn_centers/learn_width/learn_temperature to out-dir naming
+
+- Add bigger WGAN config (hidden_dim=512, n_blocks=6)
+
+- Add data-integrity guards against silent NaN/Inf propagation and races
+
+
+### Fixed
+
+- Fix column names to match actual parquet schema
+
+- Fix installed torch version to be compatible with cuda drivers
+
+- Fix miniCaloSim link in README
+
+- Fix giant.analysis import after Phase 2 dataset API changes
+
+- Fix silent failure modes surfaced by extensive code review
+
+- Fix ruff, ty, and pytest failures; apply ruff format
+
+- Clamp n_sec classification label to K_MAX
+
+- Fix rollout edep mismatch and add truth overlay to Tier 4 observables
+
+- Fix crashes in physical-property conditioning edge cases
+
+- Fix router experts silently ignoring --hidden-dim/--n-blocks
+
+- Fix conditioning="physical" so it can actually generalize past training vocab
+
+- Fix training-loop checkpoint/resume and WGAN bugs
+
+- Fix stale-partial reuse and n_chunks mismatch in analysis condor pipeline
+
+- Fix CLI/tooling robustness gaps and dedupe the Conditioning enum
+
+- Fix test_write_submit_requires_synced_venv for active-venv resolution
+
+
+### Removed
+
+- Remove scripts/train.py in favor of the giant train CLI
+
+- Drop orphaned child tracks instead of nulling secondary targets
+
+
+### Changed
+
+- Initial commit: giant surrogate model with two-phase roadmap in README
+
+- Implement Phase 1: full data pipeline, model, training, and config support
+
+- Handle material column as string type
+
+- Rename pre_energy/post_energy columns to pre_E/post_E
+
+- Rename direction columns from pre_dir_x/y/z to pre_dx/dy/dz
+
+- Batch StreamingStepsDataset internally instead of per-row collate
+
+- Dedup training pipeline, add seeding/resume and per-epoch metrics logging
+
+- Split torch into cpu/cuda extras and pin dependency version bounds
+
+- Apply ruff format and document lint/type tooling in CLAUDE.md
+
+- Update README to match current architecture and tooling
+
+- Make sampler step count configurable for validation
+
+- Calibrate auto batch size separately for inference vs training
+
+- Skip rows with unknown PDG codes during predict
+
+- Buffer predict rows across row-group boundaries before inference
+
+- Export plots for knowledge base
+
+- Rework validation notebook with markdown sections and Tier 4 plots
+
+- Allow steps_to_parquet.py to accept multiple ROOT input files
+
+- Encode edep/secondary/post energy as a conservation-constrained simplex
+
+- Expose dataset/conversion scripts as uv entry points
+
+- Restrict holdout overlap check to holdout vs dev/full only
+
+- Route predict output to UUID-named parquet with YAML reference sidecar
+
+- Implement Phase 2: secondary particle prediction
+
+- Unify dataset/tooling scripts into a single `dwarf` Typer CLI
+
+- Fold --to/--gen dataset-versioning flags into the dwarf CLI
+
+- Prefix default train output dir with current date
+
+- Color-code dwarf status output by tree level
+
+- Show VERSIONS.md reason extracts in dwarf status
+
+- Wire up predict CLI to load and run the Stage-2 sec_decoder
+
+- Wire up n_sec/species/energy-fraction validation for Stage 2
+
+- Detach Stage-2 type-embedding target to stop self-referential collapse
+
+- Weight Stage-2 secondary loss equally between direction and type-embedding dims
+
+- Recalibrate batch-size estimate for the post-Phase-2 model size
+
+- Update CLAUDE.md and README for the implemented Phase 2 model
+
+- Error on missing secondary lists instead of silently zeroing Stage-2 targets
+
+- Derive a unique per-job seed for minicalosim shard generation
+
+- Rescale secondary energies to exactly consume the e_sec budget
+
+- Support --energy-gev in dwarf make-root for the new minicalosim energy arg
+
+- Stream giant rollout output instead of buffering the whole run
+
+- Scale auto batch-size estimate by MoE expert count during training
+
+- Rewrite analysis module as a lean, fully-streaming pipeline
+
+- Reimplement rollout-vs-truth comparison on the streaming analysis module
+
+- Condition on material/particle physical properties instead of learned embeddings
+
+- Ignore the scratchpad working directory
+
+- Quote the on: key in the CI workflow
+
+- Split CI lint stage into parallel jobs
+
+- Rewrite analysis as streaming rollout-vs-reference plotting pipeline
+
+- Analyze: drive prep/submit from the rollout YAML sidecar
+
+- Analyze: show model/training params on rendered figures
+
+- Deps: install plotstyle from git.larsbogner.de package index
+
+- Analyze: drop stale ty:ignore on plotstyle import
+
+- Test: replace prep(**_CTX) splat with a typed _prep helper
+
+- Analyze: add MoE router gating/share diagnostic plots
+
+- Chore: remove stray CUDA sanity script and stale Phase 2 planning doc
+
+- Docs: document compute environment, WGAN/MoE status, and condor-gpu-train-rollout
+
+- Analyze: normalize pdg dtype in open_side to fix rollout/reference concat
+
+- Analyze: chunk per-plot aggregation across HTCondor jobs
+
+- Analyze: expose bin/pdg options on `analyze submit`
+
+- Analyze: estimate per-job HTCondor walltime from chunk row count
+
+- Analyze: run condor compute jobs via .venv/bin/giant, not uv run
+
+- Analyze: default condor docker image to alma9-gridjob
+
+- Analyze: raise default condor job memory request to 8192 MB
+
+- Analyze: recalibrate condor walltime model from real cluster timings
+
+- Transforms: pad legacy cond normalizers for pre-physical-conditioning checkpoints
+
+- Analyze: default run directory to <repo>/analysis_runs, gitignored
+
+- Docs: record first MoE router rollout benchmark result in the roadmap
+
+- Router: seed EnergyRouter centers from data quantiles instead of a fixed linspace
+
+- Docs: note the EnergyRouter centers_init fix in the roadmap
+
+- Analyze: thread full model/training/rollout/dataset params to plots
+
+- Ci: share one uv sync across jobs, gate tests on lint+type-check, sync tag/version on release tags
+
+- Ci: replace unsupported artifact sharing with a bind-mounted uv cache
+
+- Ci: stop setup-uv from overriding UV_CACHE_DIR
+
+- Ci: re-pin UV_CACHE_DIR after setup-uv, which exports its own value regardless of enable-cache
+
+- Ci: set UV_LINK_MODE=copy to silence the cross-filesystem hardlink warning
+
+- Log batch-level metrics to W&B, not just per-epoch summaries
+
+- Log router health, WGAN grad-norm split, n_sec accuracy, GPU/throughput to W&B
+
+- Persist global_step across --resume so W&B step stays monotonic
+
+- Timestamp default checkpoint dir to avoid W&B run-id collisions
+
+- Skip empty-slice mean/std in sec phys validation print
+
+- Speed up giant train's setup stage
+
+- Speed up _WelfordAccumulator's per-chunk update
+
+- Make default checkpoint out_dir name reflect only non-default hyperparams
+
+- Cache giant train's setup stage in a sidecar file
+
+- Pass --seed through to the train/val event split
+
+- Offset event_id per file to avoid cross-file collisions
+
+- Store a quantile grid instead of a raw reservoir sample in the setup cache
+
+- Scope wandb run config to only-active hyperparameters
+
+- Resolve giant condor wrapper from the active venv, not a hardcoded path
+
+- Bump version to 0.2.0
@@ -7,26 +7,32 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
 ```bash
 uv sync --extra cpu                # install dependencies with CPU-only torch (standard/default)
 uv sync --extra cuda               # install dependencies with CUDA 11.8 torch
-uv sync --extra cpu --extra dev    # add dev extras (pytest, etc.)
+uv sync --extra cpu --extra dev    # add dev extras (pytest, ruff, ty, bump-my-version, git-cliff, + all runtime extras)
 uv sync --extra cpu --extra geometry  # add scikit-learn for the geometry oracle (giant rollout)
+uv sync --extra cpu --extra workflow  # add b2luigi for `giant workflow` pipeline orchestration
 pytest               # run tests
 giant new-run --hidden-dim 512 --lr 3e-4        # scaffold a config.toml + run dir ahead of training
-giant train path/to/steps.parquet --mode flow   # train (flow matching)
-giant train path/to/steps.parquet --mode ddpm   # train (DDPM baseline)
-giant train path/to/steps.parquet --mode wgan   # train (WGAN-GP, single-pass eval; implemented, not yet tested)
-giant train path/to/steps.parquet --router --router-type energy  # MoE routing trunk (implemented; first rollout benchmark failed with lambda_balance=0, retrain needed — see Roadmap)
+giant train path/to/steps.parquet               # train (defaults: stage 1 flow, stage 2 wgan + autoregressive)
+giant train path/to/steps.parquet --mode flow   # set both stages' generative objective at once
+giant train path/to/steps.parquet --stage1-generator flow --stage2-generator wgan   # per-stage override
+giant train path/to/steps.parquet --router --router-type energy  # MoE routing trunk (see Roadmap for status)
+giant model summary --config config.toml         # build-only: parameter counts + which config keys actually bite
 giant predict path/to/steps.parquet --checkpoint ckpt/best.pt   # per-step predictions
 giant rollout path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl  # full showers
-giant analyze submit rollout.yaml --accounting-group cms   # parallel rollout-vs-reference analysis on HTCondor
-giant analyze render <run_dir> --gallery                   # render PDFs + HTML gallery (run_dir from prep/submit)
+giant workflow run spec.toml --batch --workers 20          # whole pipeline (cache-warm -> train -> rollout -> analysis)
+giant analyze prep rollout.yaml --chunks 32                # lay out an analysis run dir (compute jobs come from the workflow)
+giant analyze render <run_dir> --gallery                   # render PDFs + HTML gallery (run_dir from prep)
+giant analyze metrics <train_run_dir>                      # training-progress plots from metrics.csv
 dwarf --help                                     # dataset/tooling CLI: convert, migrate, bump-gen,
                                                  # bump-schema, status, update-manifest, create-manifest,
                                                  # make-root, build-geometry-oracle, warm-cache, hparam-scan
-                                                  # (see scripts/dwarf.py)
+                                                  # (see giant/tools/dwarf.py)
 ```

 `cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x; newer torch requires newer NVIDIA drivers). Plain `uv sync` with no extra will not install torch at all; uv has no concept of a "default extra", so `--extra cpu` should always be included unless you need GPU support.

+`configs/` holds kept reference configs (`baseline.toml`, `default.toml`, the router/WGAN scan configs) — pass them with `--config`.
+
 ### Lint and type checking

 ```bash
@@ -37,60 +43,86 @@ uv run ty check .       # type check

 Part of the `dev` extra. Run these periodically (not just at commit time) to catch drift early.

+### Release tooling
+
+Merges to `master` auto-bump the patch version, tag, and update `CHANGELOG.md` via the Gitea workflow in `.gitea/workflows/ci.yml` (bump-my-version + git-cliff). Don't hand-edit the version in `pyproject.toml` or write changelog entries by hand.
+
 ## Compute environment

 Work on this repo happens across three kinds of machine:

 - **Local dev machines** (laptop + desktop, identical): repo at `~/Programming/giant`, no access to `/ceph` — datasets, training results, and models aren't reachable here.
 - **Portal machines** (`portal1`, `deepthought`, `deepthought2`, `bms1`, `bms2`, `bms3`): repo lives under `/work`, and `/ceph` holds ROOT/parquet files and trained models. **These are shared with other users** — stay strictly within `/work/lbogner` and `/ceph/lbogner`, and keep resource usage to roughly a quarter of CPU/RAM and a single GPU so as not to disturb other users' jobs.
- **HTCondor worker nodes**: never run or SSH onto these directly — the only sanctioned path is submitting jobs through condor (`giant analyze submit`, and the in-progress remote-GPU train/rollout submission on `condor-gpu-train-rollout`). `/ceph` is available there; `/work` is only sometimes mounted, depending on the node.
+- **HTCondor worker nodes**: never run or SSH onto these directly — the only sanctioned path is `giant workflow run <spec.toml> --batch` (b2luigi, see the Workflow section), which submits and polls every job. `/ceph` is available there; `/work` is only sometimes mounted, depending on the node.

 ## Architecture

-GIANT is a conditional generative surrogate for the Geant4 step function. It replaces the stochastic physics engine: given a pre-step particle state (conditioning), it samples a post-step outcome — now including the variable-length list of secondary particles the step produces (Phase 2, see Roadmap).
+GIANT is a conditional generative surrogate for the Geant4 step function. It replaces the stochastic physics engine: given a pre-step particle state (conditioning), it samples a post-step outcome — including the variable-length list of secondary particles the step produces.

-**Data pipeline** (`giant/data/`): parquet files from miniCaloSim are loaded into numpy arrays (`loader.py`), then log-transformed and rotated into a local coordinate frame where `pre_dir = ẑ` (`transforms.py`), before being wrapped in a PyTorch `Dataset` (`dataset.py`). Train/val split is by `event_id` to avoid leaking correlated steps from the same shower.
+**Data pipeline** (`giant/data/`): parquet files from miniCaloSim are loaded into numpy arrays (`loader.py`), then log-transformed and rotated into a local coordinate frame where `pre_dir = ẑ` (`transforms.py`), before being wrapped in a PyTorch `Dataset` (`dataset.py`, streaming variant included). Train/val split is by `event_id` (`--seed`-controlled) to avoid leaking correlated steps from the same shower. Loading a directory or `.manifest` of several parquet files offsets each file's `event_id`s by a per-file stride so ids stay globally unique. `setup_cache.py` persists the pre-epoch setup scan (vocab maps, event split, process maps, normalizer stats) as a sidecar so repeated runs over the same `data` path don't rescan (`--cache-setup`/`--rebuild-setup-cache`, precomputable with `dwarf warm-cache --config ...`).

 **Stage-1 output space (9D, `giant/constants.py:LOCAL_TARGET_NAMES`):** `log_step_length`, two additive-log-ratio (ALR) coordinates `edep_logit`/`sec_logit` of a **deposit / secondary / post-energy simplex**, `post_dir` (post-scattering momentum direction, unit vector in the local frame), and `travel_dir` (direction of `post_pos - pre_pos`, unit vector in the local frame). The energy simplex decodes via softmax over `[edep_logit, sec_logit, 0]` × `pre_E` so `edep + e_sec + post_E == pre_E` holds by construction — energy conservation is architectural, not learned (see `energy_simplex_decode`). `post_pos` is not a raw target — it's reconstructed at inference as `pre_pos + step_length * world_frame(travel_dir)`, since `step_length` already encodes that displacement's magnitude and duplicating it would let the two become inconsistent.

-**Conditioning vector (15D continuous, `COND_DIM`):** pre-step position, log(pre-energy), pre-step direction, layer ID (`COND_DIM_BASE=8`) — plus, since particle/material physical-property conditioning (`model.conditioning`, see below), 7 more columns: particle `log(mass)`/`charge` (`PARTICLE_PHYS_DIM=2`, `giant/particles.py`) and material `Z_eff`/`A_eff`/`log(density)`/`log(X0)`/`log(λ_int)` (`MATERIAL_PHYS_DIM=5`, `giant/materials.py`). `n_sec` and `e_sec` are **not conditioning inputs** (that was Phase 1 / the energy-conservation PoC); the model predicts them.
+**Conditioning vector (15D continuous, `COND_DIM`):** pre-step position, log(pre-energy), pre-step direction, layer ID (`COND_DIM_BASE=8`) — plus 7 physical-property columns: particle `log(mass)`/`charge` (`PARTICLE_PHYS_DIM=2`, `giant/particles.py`) and material `Z_eff`/`A_eff`/`log(density)`/`log(X0)`/`log(λ_int)` (`MATERIAL_PHYS_DIM=5`, `giant/materials.py`). `n_sec` and `e_sec` are **not conditioning inputs** — the model predicts them. `giant/cond_layout.py` is the single source of truth for the `cond_cont`/`cond_cat` column layout shared by `giant.data.transforms`, `giant.model.encoders`, and `giant.model.routers`.

-`ConditionEncoder`/`SecondaryConditionEncoder` (`giant/model/network.py`) support two mutually exclusive `conditioning` modes, selected per-checkpoint (`model_config["conditioning"]`, defaulting to `"embedding"` for old checkpoints without the key, `"physical"` for new `giant train` runs — see `--conditioning`):
- **`"embedding"`** (original Phase 2 design): a learned `nn.Embedding` per PDG code / material name, indexed by a dataset-scoped dense vocab (`pdg_map`/`mat_map`). Memorizes the training menu.
- **`"physical"`** (default): the 7 physical-property columns above are each routed through a small MLP (`particle_mlp`/`material_mlp`) to the same `emb_dim` width the embedding tables would have produced — a drop-in replacement computable for any PDG code / material name, not just ones seen in training, which is what lets the surrogate generalize to a held-out material or species. `giant/particles.py` decodes nuclear/ion PDG codes (the `10LZZZAAAI` scheme) via the scikit-HEP `particle` package with a Z/A-digit-decode fallback for isomer codes the package's ground-state-only table misses. `giant/materials.py` ships real Geant4-11.4.1-derived `z_eff`/`a_eff`/`density`/`x0`/`lambda_int` values for every material the detector geometry actually produces; the sole exception is `G4_LYSO` (not a stock Geant4 NIST material, never actually constructed by the geometry — see the module docstring), which stays `MaterialProperties(None, ...)` and raises loudly (`MaterialPropertiesNotFilledError`) rather than silently defaulting if it's ever requested.
+`ConditionEncoder` (`giant/model/encoders.py`) configures the particle and material identity axes **independently** (`conditioning.particle` / `conditioning.material`, each a `ConditioningAxisConfig` with `type`/`emb_dim`/`n_layers`), so they may mix freely. Three per-axis modes:
+- **`"physical"`** (default): the axis's raw physical properties routed through a small MLP — computable for any PDG code / material name, which is what lets the surrogate generalize beyond the training menu. `giant/particles.py` decodes nuclear/ion PDG codes (the `10LZZZAAAI` scheme) via the scikit-HEP `particle` package with a Z/A-digit-decode fallback for isomer codes the package's ground-state-only table misses. `giant/materials.py` ships real Geant4-11.4.1-derived values for every material the detector geometry actually produces; the sole exception is `G4_LYSO` (not a stock Geant4 NIST material, never actually constructed by the geometry — see the module docstring), which stays `MaterialProperties(None, ...)` and raises loudly (`MaterialPropertiesNotFilledError`) rather than silently defaulting.
+- **`"embedding"`**: a learned `nn.Embedding` per PDG code / material name, indexed by a dataset-scoped dense vocab. Memorizes the training menu; the generalization-comparison baseline, and the only mode compatible with `stage2_model.particle_type.target = "embedding"`.
+- **`"onehot"`**: a fixed, unlearned vector over the top `emb_dim - 1` codes by training-set count plus one "other" bin. Not a reparameterization of `"embedding"` — the vocabulary cap is the real difference.

-**Model** (`giant/model/network.py`): a two-stage model, both checkpointed together.
- **Stage 1 — `DenoisingMLP`:** `ResBlock` stack with a `SinusoidalEmbedding` for the flow/diffusion time variable and a `ConditionEncoder` fusing the conditioning. Predicts the 9D primary vector field, plus an `n_sec_head` classifier over `{0..K_MAX}` (`K_MAX=15`) that runs on the condition encoding alone (no diffusion noise), callable via `predict_n_sec`.
- **Stage 2 — `SecondaryDecoder`:** a second flow-matching net (`SecondaryConditionEncoder` fuses the pre-step conditioning with the Stage-1 outcome) that generates all `K_MAX` secondary slots at once. Each slot is `(stick-breaking energy logit, local-frame direction 3D, log-mass, charge)` = `SEC_SLOT_DIM=6`, ordered by descending energy; slots beyond the predicted `n_sec` are masked. Secondary energies are a **stick-breaking partition of the `e_sec` budget** from Stage 1 (they sum to it), so the whole chain conserves energy. A secondary's mass/charge are regressed directly against a fixed physics-derived target (its ground-truth PDG code's `giant.particles.particle_mass_charge`) — not a learned/moving embedding target, so nothing needs detaching. **No snapping at inference**: the predicted (mass, charge) are used as-is as the secondary's physical identity, including for its own future conditioning if it goes on to take further steps in a rollout. A separate, reporting-only nearest-known-PDG lookup (`giant.particles.nearest_known_pdg`) is used purely to populate a nominal `pdg` label for output rows / `"embedding"`-mode fallback conditioning — it never feeds back into the model.
+`conditioning.share_stages` decides whether the two stages get one shared encoder instance or two identically-configured independent ones.

-`schedule.py` provides both a `CosineSchedule` for DDPM and the flow matching loss utilities (Lipman et al. 2022 conditional flow matching).
+**Model** (`giant/model/`, both stages checkpointed together). `network.py` is only a re-export shim now; the real code is split by concern:
+- `layers.py` — `ResBlock`/`AdaLNResBlock` + `BLOCK_REGISTRY` (conditioning-injection mechanism is selectable), `SinusoidalEmbedding`, `ContextAdapter`, `build_mlp_head`.
+- `encoders.py` — `ConditionEncoder` (above).
+- `trunks.py` — `TRUNK_REGISTRY`/`build_trunk`: everything downstream of the fused conditioning vector, as a registrable expert *body* (`resmlp` default, plus a `none` variant). `RoutedTrunk` builds `router.n_experts` instances of whichever body is named, so mixing is orthogonal to which body is mixed.
+- `routers.py` — `Router` base + `ROUTER_REGISTRY`: `energy`/`pdg`/`process`/`composed`/`none`. Soft-mixed at train time, **top-1 dispatched at eval time** (each row runs exactly one small expert), which is the actual inference-speed win. `EnergyRouter`/`PdgRouter` gate on a quantity known at inference; `ProcessRouter` runs its own small classifier (process isn't known upfront); `ComposedRouter` gates jointly over outer-product expert cells via repeated `--router-axis "type:key=val,..."`. The `--router*`/`--n-experts` CLI flags target `stage1_model.router` only; stage 2's router is config-file-only (`stage2_model.router`). `EnergyRouter` accepts `centers_init`, which `giant/pipeline.py` auto-populates from real data quantiles via a reservoir sample collected during the normalizer-fitting pass.
+- `history.py` — `HISTORY_REGISTRY`/`build_history`: `markov` (previous token only), `attention` (causal self-attention, KV-cached at inference via `init_cache`/`step`), `none`. Stage-2 autoregressive only.
+- `objectives.py` — `Objective` base + registry for `flow`/`ddpm`/`wgan`: answers in one place whether a stage needs a time embedding, is adversarial, folds the secondary type slice into its trunk output, what its trunk input is, and which loss it trains against.
+- `models.py` — the composed stage models: `Stage1Model`, `Stage2OneShot`, `Stage2Autoregressive`, `CriticModel`, all on a shared `StageModel` base.
+- `builders.py` — `build_models`/`build_critics`, assembling the above from a config dict.
+- `schedule.py` (`CosineSchedule` for DDPM + conditional-flow-matching losses), `wgan.py` (gradient penalty / critic / generator losses, Gulrajani et al. 2017), `summary.py` (`giant model summary`), `_legacy.py` (v0.2 checkpoint migration).

-**Samplers** (`giant/sample.py`): DDPM, DDIM, and flow matching (ODE integration, ~10 steps). Flow matching is the primary mode.
+**Stage 1 — primary step.** Trunk (routed or not) over the fused conditioning, plus a `SinusoidalEmbedding` of the flow/diffusion time for non-adversarial objectives, predicting the 9D vector field. An `n_sec` classifier head over `{0..k_max}` runs on the condition encoding alone; `stage2_model.n_sec.owner` decides whether it lives on stage 1 (v0.2 checkpoints) or stage 2 (default).

-**WGAN-GP mode (`--mode wgan`, implemented, not yet tested):** a throwaway fast-eval alternative to the flow/DDPM samplers above — single forward pass instead of ~10 ODE steps. Dedicated noise-conditioned generators (`WGANGenerator`/`WGANSecondaryGenerator`, `giant/model/network.py`) stand in for `DenoisingMLP`/`SecondaryDecoder`, trained against `Critic`/`SecondaryCritic` discriminators with the gradient-penalty loss in `giant/model/wgan.py` (Gulrajani et al. 2017); `sample_wgan` (`giant/sample.py`) does the single-pass draw at inference. Not yet validated against the flow-matching baseline.
+**Stage 2 — secondaries.** Conditioned on the pre-step state plus a projected stage-1 outcome (`stage2_model.context_dim`; `stage1_context` selects ground-truth vs sampled context, annealable via `ctx_p_start`/`ctx_p_end`). Two decoders (`stage2_model.decoder`):
+- **`autoregressive`** (default): one secondary at a time in descending-energy order, each token conditioned on a `HistoryEncoder` summary of prior tokens, with teacher forcing (`always`/`scheduled`/`never`, `tf_p_start`/`tf_p_end`). `n_sec.mode = "stop_token"` lets the length be emitted by the sequence itself instead of the classifier head.
+- **`one_shot`**: all `k_max` slots in one pass, masked past the predicted `n_sec` (the v0.2 behaviour).

-**MoE routing trunk (`--router`, implemented; first rollout benchmark shows the experts don't specialize — see Roadmap):** an alternative to `DenoisingMLP`'s monolithic `ResBlock` trunk — a `Router` (`giant/model/network.py`, `ROUTER_REGISTRY`/`build_router`) gates between small per-expert `ResBlock` stacks (`Expert`), soft-mixed over all experts at train time but **top-1 dispatched at eval time** (each row runs exactly one small expert), which is the actual inference-speed win. Router types gate on different conditioning axes: `EnergyRouter`/`PdgRouter` read a quantity already known at inference time, `ProcessRouter` runs its own small classifier over pre-step conditioning (since process isn't known upfront); `ComposedRouter` gates jointly over multiple axes (outer-product expert cells) via repeated `--router-axis "type:key=val,..."` flags. Config lives under `model.router` (`giant/config.py`), deep-merged one level so `router.enabled` alone doesn't drop the rest of the defaults.
+Secondary energies are a **stick-breaking partition of the `e_sec` budget** from Stage 1 (they sum to it), so the whole chain conserves energy. Particle identity is set by `stage2_model.particle_type.target`: `"onehot"` (default — categorical over the top `n_classes - 1` PDG codes by training count plus "other", with configurable `other_policy` and `class_weighting`), `"physical"` (continuous `(log-mass, charge)` regressed against `giant.particles.particle_mass_charge`), or `"embedding"` (nearest-row snap into the conditioning embedding table; requires `conditioning.particle.type = "embedding"`).

-**Validation** (`giant/validate.py`): step-level marginal comparisons.
+**Samplers** (`giant/sample.py`): DDPM, DDIM, flow matching (ODE integration, ~10 steps), and single-pass WGAN, plus the stage-2 secondary sampling loop (one-shot and autoregressive).

-**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one autoregressive `giant rollout` (for a given checkpoint) against a held-out miniCaloSim reference steps file, and produces publication-styled PDFs assembled into an HTML gallery. It exploits the fact that rollout output and a raw reference file share a world-frame physical column subset under identical names (`pre_*`/`post_*`/`edep`/`step_length`/`pdg`/`material`/`event_id`), so no ALR/local-frame decode is needed — everything is world-frame mm/MeV. Structure: `sources.py` (canonical LazyFrames + synthetic-termination-row filtering + the secondary view, which is `generation>0 & step_no==0` rollout tracks vs exploded `sec_*_list` reference columns), `reduce.py` (the streaming primitives — a single `hist1d` `group_by([group,bin]).len()` pass, per-event scalars, edep-weighted depth/transverse profiles, species share, leakage), `grouping.py`/`context.py` (fixed bin edges + energy-quantile/pdg/material group sets resolved once by `prep` into `shared.json`, so every compute job is one pass with no range scan), `catalog.py` (the declarative `PlotSpec` registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles, species/leakage, secondaries), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`). **Input is a `giant rollout` YAML sidecar** (`condor.py:load_rollout_yaml`): its `output`/`dataset` keys name the rollout parquet and the seed file (= the reference truth), and the rest of the YAML (checkpoint, geometry oracle, cutoffs) flows into each plot's gallery metadata. `prep` derives its own **run directory** next to the rollout parquet (`<...>/analysis_<id>/`) holding `shared.json`, `run_meta.json`, `reduced_partial/`, `reduced/`, `plots/`. **Compute/merge/render split:** `giant analyze submit rollout.yaml --chunks N` runs `prep` (recording the run's chunk count `N` in `run_meta.json`) then submits one HTCondor job per (plot, chunk) pair (`compute-one --id --chunk --run-dir`, polars/numpy only — no LaTeX on workers), each streaming over an `event_id`-disjoint slice (`event_id % N == chunk`) and writing a small `reduced_partial/<id>__<chunk>.json`; every `PlotSpec` (`catalog.py`) splits into a `compute_partial`/`finalize` pair so a plot's chunks can be summed/concatenated back together correctly (`chunkable=False` specs — the router diagnostics, already bounded/subsampled — always run as a single chunk regardless of `N`). The local `giant analyze render <run_dir>` first joins every plot's chunk partials into `reduced/<id>.json` (`merge_all`, a no-op join when `N=1`), then turns those into the styled PDF/gallery tree. See `giant/analysis/__init__.py`.
+**Training** (`giant/training/`): `loop.py` (epoch loop, graceful shutdown, best-checkpoint selection), `trainers.py` (`StageSpec` + per-stage flow/ddpm and WGAN-GP trainers, and the `MetricSpec` declarations that define `metrics.csv`'s columns), `stage2_inputs.py` (ground-truth stage-2 targets + teacher-forcing inputs), `metrics.py` (`MetricsCollector`: `metrics.csv`, W&B logging, progress/summary), `checkpoint.py`, `amp.py` (`train.precision = fp32|bf16` autocast), `plots.py` (`giant analyze metrics`). Per-stage `init_from`/`freeze` lets one stage be retrained against a fixed, known-good other stage while still producing a complete rollout-capable checkpoint.

-**Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower — each primary post-step becomes the next pre-step, secondaries are pushed as new tracks, and per-step `material`/`layer_id` come from a `GeometryOracle` (`giant/geometry.py`, built via `dwarf build-geometry-oracle`) that learns position → (material, layer_id) from data and flags detector escape by nearest-neighbour distance. Tracks terminate on energy cutoff, per-track max steps, escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.
+**Config** (`giant/config.py`): frozen dataclasses are the single source of truth; `DEFAULT_CONFIG` is *generated* from `GiantConfig().to_dict()` rather than hand-maintained. Blocks: `[conditioning]`, `[stage1_model]`, `[stage2_model]`, `[train]`, `[meta]`. Unknown keys are rejected on merge (with a did-you-mean suggestion), and `tests/test_config_consumed_keys.py` audits that every key is actually read somewhere.
+
+**Validation** (`giant/validate.py`): step-level marginal + KL-divergence comparisons during training (`--validate-every`).
+
+**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one or more autoregressive `giant rollout` runs against a single held-out miniCaloSim reference steps file shared by all of them, and produces publication-styled PDFs assembled into an HTML gallery — one distinctly colored series per rollout, one reference line/panel. It exploits the fact that rollout output and a raw reference file share a world-frame physical column subset under identical names (`pre_*`/`post_*`/`edep`/`step_length`/`pdg`/`material`/`event_id`), so no ALR/local-frame decode is needed — everything is world-frame mm/MeV. Structure: `sources.py` (canonical LazyFrames + `RolloutSpec`/`Side` — a rollout's opened frames + per-checkpoint diagnostic inputs — + synthetic-termination-row filtering + the secondary view, which is `generation>0 & step_no==0` rollout tracks vs exploded `sec_*_list` reference columns), `variables.py` (the per-step value expressions shared by range sizing and the plot registry), `reduce.py` (the streaming primitives — a single `hist1d` `group_by([group,bin]).len()` pass, per-event scalars, edep-weighted depth/transverse profiles, species share, leakage), `grouping.py`/`context.py` (fixed bin edges + energy-quantile/pdg/material group sets resolved once by `prep` into `shared.json` over the union of the reference and every rollout, so every compute job is one pass with no range scan), `reduced.py` (`Partial`/`Reduced` — the compact self-describing JSON a compute job emits), `catalog.py` (the declarative `PlotSpec` registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles/containment, species/leakage, secondaries, distance/confusion summaries, router and type-embedding diagnostics; `giant analyze list` prints every id), `runtime_estimate.py` (per-(plot, chunk) walltime estimates for the submit description), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`; each rollout gets a stable `ps.get_color(i)` slot by its position in `series`, the reference always draws in one fixed dashed-ink style). `Bundle.rollouts` is a name-keyed dict of `Side`, and every `compute_partial`/`finalize` builds a `Reduced.payload["series"]` dict keyed the same way, with `payload["reference"]` as the one distinguished non-rollout entry. The heatmap-shaped specs (`marginal_distance_summary`, `n_sec_confusion`) and the checkpoint-bound diagnostics (`router_gating.py`, `type_embedding_distance.py`) are inherently one-matrix/one-checkpoint per rollout, so they render as one panel per rollout instead of one line/bar per rollout.
+
+**Input is one or more `giant rollout` YAML sidecars** (`run.py:load_rollout_yamls`): each YAML's `output`/`dataset` keys name its rollout parquet and seed file (= the reference truth); every supplied YAML must resolve to the same `dataset`, checked up front with a clear error otherwise (the premise is "N candidates vs one ground truth"). Each rollout's series name comes from a repeated `--label` CLI flag, else the YAML stem (N>1), else `"rollout"` (a single YAML). `prep` creates a **run directory** (`<cwd>/analysis_runs/analysis_<id>/` by default, `--run-dir` to override) holding `shared.json`, `run_meta.json` (`RunMeta.rollouts: list[{name,path,plot_meta}]`, insertion order = CLI order = every plot's series order), `reduced_partial/`, `reduced/`, `plots/`. **Compute/merge/render split:** `giant analyze prep a.yaml [b.yaml ...] --chunks N` records `N` in `run_meta.json`, and the workflow's `AnalysisComputeTask` runs one HTCondor job per (plot, chunk) pair (`compute-one --id --chunk --run-dir`, polars/numpy only — no LaTeX on workers), each streaming over an `event_id`-disjoint slice (`event_id % N == chunk`) of the reference **and every rollout** and writing a small `reduced_partial/<id>__<chunk>.json`; every `PlotSpec` splits into a `compute_partial`/`finalize` pair so chunks can be summed/concatenated back per rollout (`chunkable=False` specs — the checkpoint-bound diagnostics, already bounded/subsampled — always run as a single chunk). The local `giant analyze render <run_dir>` first joins every plot's chunk partials into `reduced/<id>.json` (`merge_all`, a no-op join when `N=1`; `merge-one` does a single plot for debugging), then turns those into the styled PDF/gallery tree. `giant analyze metrics <train_run_dir>` is a separate, unrelated entry point: training-progress plots straight from a run's `metrics.csv`.
+
+**Workflow orchestration** (`giant/workflow/`, `giant workflow run` CLI): b2luigi is the **only sanctioned way to run a multi-step pipeline**; `giant`/`dwarf` are single-step primitives the tasks invoke. One workflow TOML (`configs/workflow_example.toml`) parameterises a whole experiment — `[workflow]`/`[condor]`/`[dataset]`/`[geometry]` plus repeated `[[train]]`/`[[rollout]]`/`[[analysis]]` tables, each cross-referenced by name — and `spec.py` parses it into frozen dataclasses, rejecting unknown keys and dangling references. Every task's output directory is `<result_dir>/<kind>/name=<name>/spec_hash=<hash>/…`, where the 8-hex `spec_hash` covers that task's resolved sub-spec **and its transitive parents**, so an edited spec re-runs exactly the affected subtree instead of silently reusing stale outputs. The DAG (`tasks.py`): `DatasetTask` (external, fails fast if `/ceph` isn't mounted) → `WarmCacheTask` / `GeometryOracleTask` → `TrainEpochTask(name, milestone)` → `TrainTask` → `RolloutTask` → `AnalysisPrepTask` → `AnalysisComputeTask(name, plot_id, chunk)` → `AnalysisRenderTask` → `WorkflowTask`. Training is fanned out into **one short GPU job per epoch** (`epochs_per_job` trades queue waits back), chained by `--resume` on the previous job's `last.pt` — the loop already handles that unchanged — and `TrainTask` republishes `best.pt`/`last.pt`/a concatenated `metrics.csv` so nothing downstream sees the fan-out. `StreamingStepsDataset.set_epoch` and `config.epoch_seed` (both applied per epoch by `training/loop.py`) derive the batch order and the global RNG state from `(seed, epoch)`, so epoch *k* is bit-identical either way — verified by diffing a chained run's `metrics.csv` against a single 3-epoch `giant train`. `AnalysisRenderTask` is always local (the only step importing plotstyle/LaTeX); `htcondor.py` holds the CPU/GPU submit settings, with the GPU requirement strings (`TARGET.ProvidesEtpCeph` + device/memory pins) ported from the `condor-gpu-train-rollout` branch. `run.py` is the script b2luigi re-executes on workers (`--spec` forwarded via `task_cmd_additional_args`, so a worker resolves the identical graph); `giant workflow run` is a thin exec of it. Needs `uv sync --extra cpu --extra workflow`.
+
+**Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower, advancing tracks breadth-first (every sweep steps all active tracks once, in `batch_size` chunks, so many tracks share each forward pass). Each primary post-step becomes the next pre-step, secondaries are pushed as new tracks, and per-step `material`/`layer_id` come from a `GeometryOracle` (`giant/geometry.py`, built via `dwarf build-geometry-oracle`) that learns position → (material, layer_id) from data and flags detector escape by nearest-neighbour distance. Tracks terminate on one of the `TERM_*` reasons in `constants.py` (energy cutoff, max steps, escape, natural end, unknown pdg, max tracks); energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction. `giant/checkpoint_io.py` is the shared checkpoint → ready-to-run-models path used by both `predict` and `rollout`.

 ## Roadmap

 **Phase 1 (done):** number of secondaries and their total energy were conditioning inputs; the model predicted only the 9D primary post-step (energy-conservation PoC).

-**Phase 2 (implemented — baseline):** the two-stage model above jointly predicts `n_sec`, the energy simplex (`e_sec` falls out of it), and each secondary's energy/direction/species, so a rollout is self-contained (no ground-truth secondary counts injected). This is the "get a baseline out" track agreed with Jan & Tobias (2026-07-07).
+**Phase 2 (done):** the two-stage model jointly predicts `n_sec`, the energy simplex (`e_sec` falls out of it), and each secondary's energy/direction/species, so a rollout is self-contained (no ground-truth secondary counts injected).

-**Physical-property conditioning (implemented):** `model.conditioning = "physical" | "embedding"` (see above) replaces the learned PDG/material embeddings with a small MLP over particle mass/charge and material Z_eff/A_eff/density/X0/λ_int, and Stage 2 predicts a secondary's mass/charge directly instead of a snapped species embedding. `"embedding"` stays available as the generalization-comparison baseline. `giant/materials.py`'s table is already filled with real values for every material the geometry produces. **Not yet done:** the actual held-out-material/species generalization comparison against the `"embedding"` baseline is unrun — the 34GB multi-material dataset at the repo root (6 materials, 237 PDG codes including nuclear/ion codes) is the natural dataset for that experiment.
+**Physical-property conditioning (implemented, default):** `conditioning.particle.type` / `conditioning.material.type` = `physical | embedding | onehot`. **Not yet done:** the actual held-out-material/species generalization comparison against the `"embedding"` baseline is unrun — the 34GB multi-material dataset (6 materials, 237 PDG codes including nuclear/ion codes) is the natural dataset for that experiment.

-**Faster-eval architectures (implemented, validation in progress):** both tracks below target a ~10× native-Geant4 eval budget and are now wired into `giant train`/`giant/model/network.py`, but neither has a validated result yet — treat both as unproven until the corresponding analysis run says otherwise:
- **WGAN-GP** (`--mode wgan`, see Architecture above): implemented, **not yet tested** — no rollout-vs-reference analysis run against it yet.
- **MoE routing trunk** (`--router`, see Architecture above): implemented, **first rollout benchmark done (2026-07-22), result: needs retraining with a different router config, not abandoned.** A 10-expert `EnergyRouter` run (`n_experts=10`, `temperature=0.5`, `learn_centers=true`, **`lambda_balance=0.0`**, only 20 fine-tuning epochs resumed from a non-routed checkpoint) diverged badly from Geant4 on step granularity, secondary species, and shower shape, despite roughly matching bulk total deposited energy. The `router_gating` diagnostic plot points at the likely cause: the ten experts overlap heavily across ~5 decades of pre-step energy instead of partitioning it — even the top-energy expert only reaches ~60–65% gate weight at the highest energies plotted — so eval-time top-1 (Voronoi) dispatch is choosing among near-ties rather than real specialists. Two contributors were identified: the missing load-balancing loss (`lambda_balance=0.0`), and `EnergyRouter`'s center init (`torch.linspace(-2, 2, n_experts)`) assuming a roughly uniform z-normalized energy distribution, which real energy spectra don't match. **Fixed (2026-07-27):** `EnergyRouter` now accepts an optional `centers_init` (backward compatible — omitting it keeps the old linspace), and `giant train` auto-populates it from real data quantiles via a reservoir sample collected during the existing normalizer-fitting pass in `giant/pipeline.py` (no extra file scan), for `--router-type energy` only. The routing *strategy* itself may still be sound, but the specific benchmarked config wasn't. **Next step before further evaluation: retrain with `lambda_balance > 0` and the new quantile-seeded centers (and consider more epochs / a from-scratch run rather than a short fine-tune), then re-check whether `router_gating` sharpens up.** Full writeup: `/home/lars/knowledge-base/experiments/giant-router-energy-rollout-validation.md`.
+**v0.3.0 — Stage-2 autoregressive redesign (implemented, released; on `master` since 2026-08-13):** motivated by the 2026-08-03 WGAN rollout benchmark, which failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). Stage 2 became autoregressive in descending-energy order with teacher forcing, and the particle-type representation went back to **categorical** (`particle_type.target = "onehot"`), reversing the 2026-07-17 continuous `(log-mass, charge)` target. The config break (`[conditioning]`/`[stage1_model]`/`[stage2_model]`/`[train]` replacing the flat `train.mode` + `[model]`) makes per-stage generators, stage-2-only training, and one-shot-vs-autoregressive comparison all expressible, and the `network.py` refactor into composable parts (encoder × trunk × objective) also makes routed WGAN work for the first time.

-A sampling-calorimeter (multi-material) dataset is still a planned future direction, not yet built. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`).
+v0.2 configs and checkpoints are auto-migrated (`config.migrate_config`, `model._legacy._migrate_legacy_model_config`, both drawing on shared facts in `giant/_migration.py`). **v0.2 checkpoint-loading support has no expiry decided yet**: `/ceph` still holds pre-v0.3.0 checkpoints and analysis runs referencing them, so don't delete or substantially alter either migration function or `tests/legacy/network_v02_snapshot.py` (the frozen v0.2 snapshot they're tested against) without an explicit decision to do so first.

-**v0.3.0 — Stage-2 autoregressive redesign (designed, not implemented; branch `v0.3.0-stage2-autoregressive`):** the 2026-08-03 WGAN rollout benchmark failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). The agreed response pivots Stage 2 to **autoregressive generation** in descending-energy order with teacher forcing, and switches the particle-type representation back to **categorical** (top N−1 by training-set count + an "other" bucket), reversing the 2026-07-17 continuous `(log-mass, charge)` target. This requires a config break: `[conditioning]` / `[stage1_model]` / `[stage2_model]` / `[train]` blocks replace the single global `train.mode` + `[model]`, so per-stage generators (`stage1 = flow` + `stage2 = wgan`), stage-2-only training, and one-shot-vs-autoregressive comparison are all expressible. `network.py` is refactored from ten permutation classes into composable parts (encoder × trunk × objective), which also makes routed WGAN work for the first time. **Full design contract, with every config option documented: `docs/v0.3.0-design.md` — read it before touching `giant/config.py` or `giant/model/network.py`.**
+**Faster-eval architectures — both implemented, neither validated.** Target is a ~10× native-Geant4 eval budget; no eval-latency number exists for any configuration yet, so that budget is unverified across the board.
+- **WGAN-GP** (`--stage2-generator wgan`, now the stage-2 default): first rollout benchmark 2026-08-03 failed with secondary-species mode collapse — the failure v0.3.0 was designed to address. **No post-v0.3.0 benchmark has been run.** Writeup: `/home/lars/knowledge-base/experiments/giant-wgan-physical-rollout-validation.md`.
+- **MoE routing trunk** (`--router`): first rollout benchmark 2026-07-22 diverged badly from Geant4 on step granularity, secondary species, and shower shape, despite roughly matching bulk total deposited energy. Cause identified as a bad config, not a bad idea: `lambda_balance=0.0` (no load-balancing loss) plus `EnergyRouter`'s `torch.linspace(-2, 2, n_experts)` center init assuming a roughly uniform z-normalized energy distribution — so the ten experts overlapped across ~5 decades of energy instead of partitioning it, and eval-time top-1 dispatch chose among near-ties rather than real specialists. Both prerequisites are fixed in code (quantile-seeded `centers_init` from `pipeline.py`, `lambda_balance` exposed). **Next step: retrain with `lambda_balance > 0` and quantile-seeded centers (consider a from-scratch run rather than a short fine-tune), then re-check whether `router_gating` sharpens up.** Writeup: `/home/lars/knowledge-base/experiments/giant-router-energy-rollout-validation.md`.

-**Condor-submitted GPU training/rollout (in progress, `condor-gpu-train-rollout` branch, not yet merged):** moves `giant train`/`giant rollout` off the shared portal GPU dev machines (see Compute environment) onto remote-GPU HTCondor submission on TOpAS/NEMO2 (`giant/condor.py`). Partway between "needs major features" and feature-complete — not ready to merge yet.
+A sampling-calorimeter (multi-material) dataset track is still open and unblocked, not yet started. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`).
+
+**Condor-submitted GPU training/rollout (`condor-gpu-train-rollout` branch, superseded):** its goal — moving `giant train`/`giant rollout` off the shared portal GPU dev machines onto remote-GPU HTCondor submission — is now met by the b2luigi workflow above. Its `train-submit`/`rollout-submit` commands are deliberately **not** ported and must not be revived when that branch is eventually merged; the only part that survived is `_gpu_requirements`, which moved into `giant/workflow/htcondor.py`.
@@ -1,20 +1,30 @@
 # giant

-**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate — a play on *Geant4* and the step function being the computationally heaviest part of the simulation.
+**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate.

-Conditional generative surrogate for the Geant4 step function. Given a pre-step particle state, the model samples a physically plausible post-step outcome — the primary's continuation plus the variable-length list of secondary particles it produces — replacing the stochastic Geant4 physics engine with a trained generative model. A trained checkpoint autoregressively rolls out full showers, stepping each primary and pushing secondaries as new tracks.
+A conditional generative model that replaces the Geant4 step function: given a pre-step particle state it samples a post-step outcome — the primary's continuation plus its secondary particles — and autoregressively rolls that out into full showers. Trained entirely from parquet dumps of the miniCaloSim steps tree; no Geant4 runtime dependency.

-Training is driven entirely from parquet files of the miniCaloSim steps tree. No Geant4 runtime dependency.
+## Quick start
+
+```bash
+uv sync --extra cpu               # install deps (CPU torch; use --extra cuda for GPU)
+
+giant new-run --hidden-dim 512 --lr 3e-4          # scaffold config.toml + run dir
+giant model summary --config config.toml          # parameter counts + which config keys actually bite
+giant train path/to/steps.parquet                 # train (flow + wgan by default)
+giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
+
+dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl   # needed for rollout
+giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
+```
+
+Every command takes `--help` for the full flag list, and `--config config.toml` for anything not exposed as a flag.

 ## Architecture

-A **two-stage model**, both stages checkpointed together, with a choice of generative mode per stage (`--mode`):
+A **two-stage model**, checkpointed together. Either stage's outcome can be produced by one of three interchangeable generative objectives (`--stage1-generator`/`--stage2-generator`, or `--mode` to set both at once): `flow` (conditional flow matching, ODE-sampled in ~10 steps), `ddpm` (denoising diffusion), or `wgan` (single-pass WGAN-GP generator/critic).

- **`flow`** (default) — conditional flow matching (Lipman et al. 2022): an MLP learns a vector field mapping noise → step outcomes, sampled via ODE integration in ~10 steps.
- **`ddpm`** — a standard denoising diffusion baseline for comparison (`giant/model/schedule.py:CosineSchedule`).
- **`wgan`** — a single-pass Wasserstein-GAN-GP generator/critic (`giant/model/wgan.py`), trading iterative sampling for one forward pass; implemented, not yet validated against the flow-matching baseline.
-
-**Stage 1 — primary (9D, `giant/constants.py:LOCAL_TARGET_NAMES`):**
+**Stage 1 — primary step.** Predicts the 9D post-step outcome (`giant/constants.py:LOCAL_TARGET_NAMES`) from the pre-step conditioning:

 | Index | Variable | Encoding |
 |-------|----------|----------|
@@ -23,36 +33,29 @@ A **two-stage model**, both stages checkpointed together, with a choice of gener
 | 3–5 | `post_dir` in local frame | unit vector |
 | 6–8 | `travel_dir` (`post_pos − pre_pos`) in local frame | unit vector |

-The two energy logits decode via softmax over `[edep_logit, sec_logit, 0]` × `pre_E`, so `edep + e_sec + post_E == pre_E` exactly — **energy conservation is built into the parametrization**, not left to the loss (`energy_simplex_decode`). Stage 1 also has a classifier head (`predict_n_sec`) predicting the number of secondaries `n_sec ∈ {0..K_MAX}` (`K_MAX = 15`) from the conditioning alone, no diffusion noise involved.
+- Energy logits decode via softmax over `[edep_logit, sec_logit, 0]` × `pre_E`, so `edep + e_sec + post_E == pre_E` exactly — conservation is architectural, not learned.
+- `post_dir`/`travel_dir` live in the frame where `pre_dir = ẑ`. `post_pos` isn't a target — it's reconstructed as `pre_pos + step_length * world_frame(travel_dir)`.

-Both `post_dir` and `travel_dir` are expressed in the coordinate frame where `pre_dir = ẑ`, making the scattering distribution nearly azimuthally symmetric. `post_pos` itself is not a raw target — it's reconstructed at inference as `pre_pos + step_length * world_frame(travel_dir)`, so the two stay consistent by construction instead of being learned (and potentially diverging) independently.
+**Stage 2 — secondaries.** Conditioned on the pre-step state and Stage 1's outcome, it generates the variable-length list of secondary particles. Two decoding strategies (`--stage2-decoder`):

-**Stage 2 — secondaries (`SecondaryDecoder`):** conditioned on the pre-step state *and* the Stage-1 outcome, a second net generates all `K_MAX` secondary slots at once — `(stick-breaking energy logit, local-frame direction, log-mass, charge)` per slot, ordered by descending energy; slots beyond the predicted `n_sec` are masked. Secondary energies are a stick-breaking partition of the `e_sec` budget from Stage 1, so the whole chain conserves energy. A secondary's mass/charge are regressed directly against its ground-truth PDG code's physical values (`giant.particles.particle_mass_charge`) and used as-is at inference — including for its own conditioning if it takes further steps in a rollout. No snapping to a known PDG code happens in the model path; `giant.particles.nearest_known_pdg` is a reporting-only lookup used to populate a nominal `pdg` label on output rows.
+- `autoregressive` — emits secondaries one at a time in descending-energy order, each token conditioned on a running history of prior tokens (`markov`: previous token only, or `attention`: causal self-attention, KV-cached at inference)
+- `one_shot` — all `K_MAX` slots generated in a single forward pass, masked past the predicted `n_sec`

-**Conditioning (`--conditioning`, per-checkpoint):** pre-step position, log(pre-energy), pre-step direction, layer ID, plus particle/material physical properties — mass/charge (`giant/particles.py`) and Z_eff/A_eff/density/X0/λ_int (`giant/materials.py`). Two mutually exclusive modes:
+Either way, secondary energies stick-break the `e_sec` budget handed down from Stage 1, so the full chain conserves energy. A secondary's particle identity is represented as `onehot` (categorical, top-N PDG codes + "other"), `physical` (continuous log-mass/charge), or `embedding` (nearest-neighbour lookup).

- **`physical`** (default) — the physical-property columns are routed through small MLPs, computable for any PDG code / material, letting the surrogate generalize to species/materials outside the training menu.
- **`embedding`** — the original design: a learned `nn.Embedding` per PDG code / material, kept as a generalization-comparison baseline (memorizes the training menu).
+**Conditioning.** Pre-step position/energy/direction/layer, plus particle mass/charge and material Z_eff/A_eff/density/X0/λ_int, encoded the same three ways as particle identity above. The particle and material axes are configured independently (`conditioning.particle.type` / `conditioning.material.type`; `--conditioning` sets both at once) and may mix — the `physical` representation generalizes to species/materials outside the training menu since it's computed rather than looked up. `n_sec`/`e_sec` are always model outputs, never conditioning inputs.

-`n_sec` and `e_sec` are model outputs, not conditioning inputs — a rollout is self-contained and never injects ground truth.
-
-**Mixture-of-experts routing (`--router`, opt-in):** `giant/model/network.py` also implements a pluggable `Router` contract (`ROUTER_REGISTRY`: `energy`, `pdg`, `process`, plus a `composed` router combining several axes) that splits `DenoisingMLP`/`SecondaryDecoder` into per-expert trunks, soft-gated in training and top-1 dispatched at eval. Implemented; first rollout benchmark needs a retrain with a load-balancing loss and better-seeded router centers (see Roadmap). See `--router-type`/`--n-experts`/`--router-axis` on `giant train`/`giant new-run`.
-
-## Roadmap
-
-**Phase 1 (done):** `n_sec` and total secondary energy `e_sec` were conditioning inputs; the model predicted only the 9D primary post-step (energy-conservation PoC).
-
-**Phase 2 (implemented — baseline):** the two-stage model above predicts `n_sec` and each secondary's energy, direction, and species jointly with the primary, so a shower rollout is fully self-contained.
-
-**Physical-property conditioning (implemented):** replaces learned PDG/material embeddings with physical-property MLPs (see above); Stage 2 predicts a secondary's mass/charge directly instead of a snapped species embedding. Not yet done: the held-out-material/species generalization comparison against the `embedding` baseline — the natural dataset for that is the 34GB multi-material dataset at the repo root (6 materials, 237 PDG codes).
-
-**Faster-eval architectures (implemented, validation in progress):** both target a ~10× native-Geant4 eval budget. WGAN-GP (`--mode wgan`) has no rollout-vs-reference analysis run against it yet. The MoE router (`--router`) had its first rollout benchmark diverge from Geant4 despite matching bulk deposited energy — the experts weren't specializing (near-uniform gating), traced to a missing load-balance loss and a center-init that didn't match the real energy distribution; both are now fixable via `lambda_balance > 0` and quantile-seeded router centers, but a re-run to confirm hasn't happened yet.
-
-A multi-material sampling-calorimeter dataset is a planned future direction, not yet built.
+**MoE routing** (`--router`): a pluggable `Router` (`energy`/`pdg`/`process`/`composed` axes) top-1-dispatches each row to one of several small expert trunks at eval time, instead of running one monolithic trunk. The CLI flags configure Stage 1's router; Stage 2 has its own `stage2_model.router` block, config-file only.

 ## Data

-Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from a ROOT file via `dwarf convert`. Each row is one Geant4 step. Train/val split is by `event_id` (not row shuffle, and `--seed`-controlled) to avoid leaking correlated steps from the same shower.
+- Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from ROOT via `dwarf convert`. One row = one Geant4 step.
+- **Conditioning (pre-step) columns:** `event_id`, `pdg`, `pre_x`/`pre_y`/`pre_z`, `pre_E`, `pre_dx`/`pre_dy`/`pre_dz` (direction), `material`, `layer_id`.
+- **Primary outcome (post-step) columns:** `post_x`/`post_y`/`post_z`, `post_E`, `post_dx`/`post_dy`/`post_dz`, `step_length`, `edep` (energy deposited in this step), `e_sec` (total energy carried off by secondaries), `child_track_ids` (its length gives `n_sec`).
+- **Secondary columns**, one variable-length list per step: `sec_pdg_list`, `sec_E_list`, `sec_dx_list`/`sec_dy_list`/`sec_dz_list` — padded/truncated to `K_MAX` (15) slots on load, ordered by descending energy.
+- **Optional:** `process` — the physics process that produced the step (e.g. `compt`, `phot`, `eBrem`); a post-step label used only as classifier supervision (`ProcessRouter`), never as conditioning.
+- Train/val split is by `event_id` (`--seed`-controlled), not row shuffle, so correlated steps from the same shower never leak across the split.
+- Loading a directory or `.manifest` of multiple parquet files (each one Geant4 job, `event_id` restarting from 0) offsets each file's `event_id`s by a fixed per-file stride so ids stay globally unique across files.

 ## Project structure

@@ -62,31 +65,61 @@ giant/
 │   ├── data/
 │   │   ├── loader.py       # parquet → numpy arrays (incl. streaming/chunked reads)
 │   │   ├── transforms.py   # log transforms, local-frame rotation, energy simplex, secondary encode/decode
-│   │   └── dataset.py      # StepsDataset / StreamingStepsDataset (PyTorch)
+│   │   ├── dataset.py      # StepsDataset / StreamingStepsDataset (PyTorch)
+│   │   └── setup_cache.py  # sidecar cache for the pre-epoch setup scan (vocab/split/normalizers)
 │   ├── model/
-│   │   ├── network.py      # ConditionEncoder, DenoisingMLP, SecondaryDecoder, Router/MoE, WGAN generator/critic
+│   │   ├── models.py       # Stage1Model, Stage2OneShot, Stage2Autoregressive, CriticModel
+│   │   ├── builders.py     # build_models / build_critics — config dict → assembled stage models
+│   │   ├── encoders.py     # ConditionEncoder (physical / embedding / onehot, per axis)
+│   │   ├── layers.py       # ResBlock/AdaLNResBlock registry, SinusoidalEmbedding, MLP heads
+│   │   ├── trunks.py       # trunk registry (resmlp, none) + RoutedTrunk (MoE expert bodies)
+│   │   ├── routers.py      # Router registry: energy / pdg / process / composed / none
+│   │   ├── history.py      # stage-2 AR history encoders: markov / attention (KV-cached) / none
+│   │   ├── objectives.py   # flow / ddpm / wgan objective registry
 │   │   ├── schedule.py     # CosineSchedule (DDPM) and flow matching utilities
-│   │   └── wgan.py         # WGAN-GP gradient penalty / critic / generator losses
+│   │   ├── wgan.py         # WGAN-GP gradient penalty / critic / generator losses
+│   │   ├── summary.py      # build-only introspection behind `giant model summary`
+│   │   ├── _legacy.py      # v0.2 checkpoint model_config/state-dict migration
+│   │   └── network.py      # re-export shim over all of the above
 │   ├── constants.py        # output/conditioning dims, K_MAX, secondary slot layout, schema keys
-│   ├── particles.py        # PDG → (mass, charge) decode, incl. nuclear/ion codes; nearest-known-PDG lookup
+│   ├── cond_layout.py      # single source of truth for the cond_cont/cond_cat column layout
+│   ├── particles.py        # PDG → (mass, charge) decode, incl. nuclear/ion codes; onehot/embedding secondary-identity decode
 │   ├── materials.py        # material name → (Z_eff, A_eff, density, X0, λ_int)
 │   ├── config.py           # default hyperparameters, TOML config merging, device autodetect
 │   ├── pipeline.py         # builds datasets/normalizers and kicks off a training run (with setup-stage caching)
-│   ├── train.py            # two-stage training loop, checkpointing, graceful shutdown, W&B logging
+│   ├── training/           # two-stage training: loop, per-stage trainers, metrics, checkpointing
+│   │   ├── loop.py         #   epoch loop, graceful shutdown, best-checkpoint selection
+│   │   ├── trainers.py     #   StageSpec + flow/ddpm and WGAN-GP per-stage trainers
+│   │   ├── stage2_inputs.py#   ground-truth stage-2 targets + autoregressive/teacher-forcing inputs
+│   │   ├── metrics.py      #   MetricsCollector: metrics.csv columns, W&B logging, progress/summary
+│   │   ├── amp.py          #   bf16 autocast (`train.precision`)
+│   │   ├── plots.py        #   training-progress plots (`giant analyze metrics`)
+│   │   └── checkpoint.py   #   checkpoint assembly/restore (format unchanged since v0.2)
 │   ├── sample.py           # DDPM / DDIM / flow matching / WGAN samplers + secondary sampling
+│   ├── checkpoint_io.py    # checkpoint → ready-to-run models/normalizers (predict + rollout)
 │   ├── geometry.py         # GeometryOracle: position → (material, layer_id, escaped) for rollout
 │   ├── rollout.py          # autoregressive shower rollout driver
 │   ├── validate.py         # step-level marginal + KL-divergence validation
+│   ├── _migration.py       # shared v0.2 → v0.3 facts used by both migration surfaces
 │   ├── analysis/           # rollout-vs-reference analysis pipeline (see `giant analyze` below)
 │   │   ├── sources.py      #   canonical LazyFrames + secondary view
+│   │   ├── variables.py    #   per-step value expressions shared by range sizing and the catalog
 │   │   ├── reduce.py       #   streaming reduction primitives (hist1d, per-event scalars, profiles, ...)
 │   │   ├── grouping.py     #   fixed bin edges + energy/pdg/material group sets
 │   │   ├── context.py      #   resolves grouping into `shared.json` once per run
-│   │   ├── catalog.py      #   declarative PlotSpec registry
-│   │   ├── condor.py       #   prep / compute-one / submit-description plumbing
+│   │   ├── reduced.py      #   Partial/Reduced — the compact JSON a compute job emits
+│   │   ├── catalog.py      #   declarative PlotSpec registry (`giant analyze list`)
+│   │   ├── router_gating.py / type_embedding_distance.py   # checkpoint-bound diagnostics
+│   │   ├── runtime_estimate.py  #   per-(plot, chunk) walltime estimates for the job requests
+│   │   ├── run.py          #   prep / compute-one / merge plumbing
 │   │   └── render.py       #   PDFs + HTML gallery (only module importing plotstyle/LaTeX)
-│   └── cli.py              # `giant train` / `new-run` / `predict` / `rollout` / `analyze` Typer app
-├── scripts/                            # dataset/tooling logic, unified under the `dwarf` CLI (`dwarf --help`)
+│   ├── workflow/           # b2luigi pipeline orchestration (`giant workflow run spec.toml`)
+│   │   ├── spec.py         #   workflow TOML -> frozen dataclasses, validation, per-task spec hashes
+│   │   ├── htcondor.py     #   CPU/GPU submit settings (docker image, +RemoteJob, GPU requirements)
+│   │   ├── tasks.py        #   the task graph: cache-warm -> train (one job/epoch) -> rollout -> analysis
+│   │   └── run.py          #   the script b2luigi re-executes on every worker
+│   └── cli.py              # `giant train` / `new-run` / `model summary` / `predict` / `rollout` / `analyze` / `workflow`
+├── giant/tools/                        # dataset/tooling logic, unified under the `dwarf` CLI (`dwarf --help`)
 │   ├── dwarf.py                        # Typer app: convert, migrate, bump-gen, bump-schema, status,
 │   │                                    #   update-manifest, create-manifest, make-root,
 │   │                                    #   build-geometry-oracle, warm-cache, hparam-scan
@@ -106,39 +139,77 @@ giant/
 ## Setup

 ```bash
-uv sync --extra cpu               # CPU-only torch (use --extra cuda for CUDA 11.8 instead)
-uv sync --extra cpu --extra dev   # add dev tools (pytest, ruff, ty)
+uv sync --extra cpu                   # CPU-only torch (use --extra cuda for CUDA 11.8 instead)
+uv sync --extra cpu --extra dev       # add dev tools (pytest, ruff, ty)
 uv sync --extra cpu --extra geometry  # add scikit-learn, for `dwarf build-geometry-oracle` / rollout
+uv sync --extra cpu --extra analysis  # matplotlib/polars/plotstyle, for `giant analyze render`
+uv sync --extra cpu --extra convert   # uproot/awkward/polars, for `dwarf convert`
+uv sync --extra cpu --extra wandb     # W&B logging (`giant train --wandb`)
+uv sync --extra cpu --extra workflow  # b2luigi, for `giant workflow run`
 ```

-`cpu` and `cuda` are mutually exclusive extras selecting the torch build (pinned to 2.3.x); plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
+The `dev` extra pulls in `convert`, `analysis`, `geometry`, `wandb` and `workflow` as well.

-## Training, prediction, and rollout
+`cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x). Plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
+
+## Training, prediction, rollout

 ```bash
-giant new-run --hidden-dim 512 --lr 3e-4 --comment "..."   # scaffold a config.toml + run dir for a new run
-giant train path/to/steps.parquet --mode flow               # train (flow matching; also --mode ddpm / wgan)
+giant new-run --hidden-dim 512 --lr 3e-4 --comment "..."   # scaffold a config.toml + run dir
+giant train path/to/steps.parquet                          # train (flow stage 1 + wgan stage 2, default)
 giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt

-# Full-shower rollout needs a geometry oracle (position → material/layer_id):
-dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl
+dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl   # position → material/layer_id
 giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
 ```

-`train`/`predict` accept a TOML config file (`--config`) and CLI overrides for hyperparameters; see `--help` on any command for the full option list. `giant train --wandb` logs per-epoch metrics (the same ones written to `metrics.csv`) to Weights & Biases; requires `uv sync --extra wandb`. A repeat `giant train` against the same dataset (e.g. a hyperparameter sweep) reuses a cached setup-stage sidecar (vocab maps, event split, normalizer stats) unless `--no-cache-setup`/`--rebuild-setup-cache`; `dwarf warm-cache` precomputes it ahead of time. `giant rollout` seeds showers from the highest-energy entry step per event, then autoregressively steps the two-stage model to completion — pushing secondaries as new tracks and looking up `material`/`layer_id` from the oracle at each step. Tracks terminate on energy cutoff, per-track max steps, detector escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.
+Useful flags on `giant train`:
+
+- `--mode {flow,ddpm,wgan}` sets both stages' objective at once; `--stage1-generator`/`--stage2-generator` override per stage
+- `--stage2-decoder {autoregressive,one_shot}` — Stage 2 decoding strategy (see Architecture)
+- `--conditioning {physical,embedding,onehot}` — conditioning representation
+- `--router` / `--router-type` / `--n-experts` / `--router-axis` — MoE routing
+- `--stage2-stage1-context {truth,sampled}` — feed Stage 2 the ground-truth or the model's own sampled Stage-1 outcome (annealable via `stage2_model.ctx_p_start`/`ctx_p_end`)
+- `--precision {fp32,bf16}` — bf16 autocast in the training loop
+- `--wandb` — log per-epoch metrics to Weights & Biases (needs `uv sync --extra wandb`); metric names are `<stage>/<split>/<metric>` plus an unprefixed run-level tail, all derived from `giant/training/trainers.py` `MetricSpec`s
+- `--no-cache-setup` / `--rebuild-setup-cache` — control the setup-stage sidecar cache (vocab maps, event split, normalizer stats); `dwarf warm-cache` precomputes it
+- `--stage1-init-from`/`--stage2-init-from` (checkpoint `.pt`) + `--stage1-freeze`/`--stage2-freeze` — load a stage's weights from another checkpoint and never update them, so the other stage can be retrained alone against a fixed, known-good one while still producing a complete, rollout-capable checkpoint
+
+Config-file-only knobs (no CLI flag — use `--config config.toml`): `stage2_model.autoregressive.teacher_forcing`/`.history`, `stage2_model.particle_type.target`/`.class_weighting`, `stage2_model.n_sec.mode`/`.owner`, `conditioning.share_stages`, `stage*_model.trunk.*` and the finer `router` knobs (`lambda_balance`, `gumbel`, `learn_width`, …). `configs/` holds kept reference configs. v0.2 flat-schema configs and checkpoints load fine (auto-migrated).
+
+`giant rollout` seeds showers from each event's highest-energy entry step, then autoregressively steps the model to completion, pushing secondaries as new tracks and looking up `material`/`layer_id` from the geometry oracle each step. Tracks terminate on energy cutoff, max steps, detector escape, or natural end; energy is deposited locally on every stop except escape, so showers conserve energy by construction.

 ## Validation and analysis

-`giant.validate.validate_marginals` runs step-level marginal and KL-divergence checks during training (`--validate-every`).
-
-For deeper rollout-vs-reference diagnostics — marginals stratified by energy/pdg/material, per-event totals, shower profiles, species share, leakage, and secondaries — `giant analyze` runs a streaming compute/render pipeline against a `giant rollout` YAML sidecar:
+- `giant.validate.validate_marginals` — step-level marginal + KL-divergence checks during training (`--validate-every`)
+- `giant analyze` — deeper rollout-vs-reference diagnostics (marginals by energy/pdg/material, per-event totals, shower profiles, species share, leakage, secondaries):

 ```bash
-giant analyze submit rollout.yaml --accounting-group cms   # prep + one HTCondor job per plot (compute only)
-giant analyze render <run_dir> --gallery                   # local: styled PDFs + HTML gallery (needs LaTeX)
+giant analyze prep rollout.yaml --chunks 8                 # lay out the run directory
+giant analyze prep a.yaml b.yaml --label flow --label wgan  # N rollouts vs one shared reference
+giant analyze render <run_dir> --gallery                   # local: merge chunks, then styled PDFs + HTML gallery (needs LaTeX)
+
+giant analyze list                                         # every catalog plot id
+giant analyze compute-one --id marginal_edep --run-dir <run_dir> --chunk 0   # what a condor job runs
+giant analyze merge-one --id marginal_edep --run-dir <run_dir>               # merge one plot's chunks (debugging)
 ```

-`<run_dir>` is derived next to the rollout parquet (`analyze prep`/`submit` print it). Compute jobs are polars/numpy only; only `render` imports plotstyle/LaTeX, so it always runs locally.
+The per-(plot, chunk) compute jobs themselves are submitted by the workflow (below), not by `giant analyze` — these commands are the single-step primitives it calls. `<run_dir>` defaults to `<cwd>/analysis_runs/analysis_<id>` (`--run-dir` overrides it; `prep` prints it). Multiple rollout YAMLs must all name the same reference (`dataset`) file; each renders as its own colored series against one reference line/panel. Compute jobs are polars/numpy only; only `render` needs LaTeX, so it always runs locally.
+
+## Workflow orchestration
+
+Multi-step pipelines run through [b2luigi](https://github.com/belle2/b2luigi) — one spec file describes a whole experiment, and every step's outputs are files on `/ceph` that are only recomputed when their spec (or an upstream one) changes:
+
+```bash
+uv sync --extra cpu --extra workflow
+giant workflow run configs/workflow_example.toml --mode dry-run   # what would run
+giant workflow run configs/workflow_example.toml --mode show-output  # where every output goes
+giant workflow run configs/workflow_example.toml --batch --workers 20  # submit to HTCondor and wait
+```
+
+The spec holds `[workflow]`/`[condor]`/`[dataset]`/`[geometry]` plus repeated `[[train]]`, `[[rollout]]` and `[[analysis]]` tables cross-referenced by name (see `configs/workflow_example.toml`). The task graph is `DatasetTask → WarmCacheTask/GeometryOracleTask → TrainEpochTask… → TrainTask → RolloutTask → AnalysisPrepTask → AnalysisComputeTask(plot, chunk) → AnalysisRenderTask`. Training is split into one short GPU job per epoch (chained by `--resume`), which schedules better on a busy farm and survives preemption; `TrainTask` then publishes one `best.pt`/`last.pt`/`metrics.csv` for everything downstream. Rendering always runs locally, since it is the only step that needs LaTeX.
+
+Separately, `giant analyze metrics <train_run_dir>` renders training-progress plots (loss/lr/accuracy/grad-norm/router/wgan/throughput) straight from a training run's `metrics.csv`.

 ## Development

@@ -0,0 +1,52 @@
+# git-cliff configuration — see https://git-cliff.org/docs/configuration
+#
+# Commit messages in this repo aren't Conventional Commits; they're plain
+# imperative summaries like "Add class-balanced secondary particle-type loss
+# (gitea #44)". Parsing here is tuned to that convention rather than to
+# feat:/fix:-style prefixes.
+
+[changelog]
+header = "# Changelog\n\n"
+body = """
+{% if version %}\
+## [{{ version | trim_start_matches(pat="v") }}] - {{ timestamp | date(format="%Y-%m-%d") }}
+{% else %}\
+## [Unreleased]
+{% endif %}\
+{% for group, commits in commits | group_by(attribute="group") %}
+### {{ group | striptags | trim | upper_first }}
+{% for commit in commits %}
+- {{ commit.message | upper_first }}
+{% endfor %}
+{% endfor %}
+"""
+trim = true
+render_always = true
+postprocessors = []
+
+[git]
+conventional_commits = false
+filter_unconventional = false
+require_conventional = false
+split_commits = false
+# Keep only the commit subject (first line), then linkify "(gitea #N)".
+commit_preprocessors = [
+    { pattern = "(?s)\n.*", replace = "" },
+    { pattern = "\\(gitea #(\\d+)\\)", replace = "[gitea #${1}](https://git.larsbogner.de/lars/giant/issues/${1})" },
+]
+protect_breaking_commits = false
+commit_parsers = [
+    { message = "^Merge ", skip = true },
+    { message = "\\[skip ci\\]", skip = true },
+    { message = "^Add", group = "<!-- 0 -->Added" },
+    { message = "^(Fix|Clamp|Clip)", group = "<!-- 1 -->Fixed" },
+    { message = "^(Remove|Drop|Deprecate)", group = "<!-- 2 -->Removed" },
+    { message = ".*", group = "<!-- 3 -->Changed" },
+]
+filter_commits = false
+link_parsers = []
+use_branch_tags = false
+topo_order = false
+topo_order_commits = true
+sort_commits = "oldest"
+recurse_submodules = false
@@ -0,0 +1,129 @@
+# GIANT reference baseline (v0.3 schema).
+#
+# The fixed comparison point every future architecture variant is measured
+# against. Chosen so that each experimental axis the roadmap cares about
+# (routed trunk, WGAN generators, attention history, shared conditioning,
+# embedding/onehot conditioning) is a *single* edit away from this file.
+#
+# Rationale for the choices below, from the runs already on record
+# (analysis_runs/ + the `giant` W&B project):
+#
+#   * flow, not wgan, for both stages. Ranking the five existing rollouts by
+#     mean Jensen-Shannon divergence against the Geant4 reference, the plain
+#     non-routed flow model wins (0.172) over the routed flow runs
+#     (0.197/0.200) and both WGAN runs (0.218/0.234) — and it beats them by
+#     ~7x on per-event total deposited energy and by 3-10x on every
+#     per-PDG marginal. WGAN stays a variant, not the reference.
+#
+#   * no router. The routed runs are not better, and soft-mixing 10 small
+#     experts costs ~10x per-pass throughput at train time (29k samples/s vs
+#     the WGAN runs' 52-116k), which is what made those runs take ~110 h for
+#     30 epochs.
+#
+#   * hidden_dim 512 / 6 blocks per stage. The best-scoring rollout so far
+#     was hidden_dim 1024, but at 4x the trunk FLOPs of 512. 512/6 sits in
+#     the same weight class as the variants it will be compared against and
+#     leaves headroom to train it properly rather than cheaply.
+#
+#   * dropout 0.0. Training set is ~5e8 steps against <1e7 parameters;
+#     capacity overfitting is not the binding constraint, and every recent
+#     run used 0.0.
+#
+# Known weak spots this baseline is expected to *exhibit* (they are the
+# reason for the comparisons, not a reason to retune this file): every model
+# on record under-produces steps per event by ~2x (rollout ~7e4 vs Geant4
+# ~1.4e5) and secondaries per event by 2-3.5x (~2-3e4 vs 7.2e4), and n_sec
+# head accuracy sits at 0.863-0.867 regardless of size or objective.
+
+[meta]
+# REQUIRED. Without it config.migrate_config reads this file as v0.2 and
+# rewrites it from V02_FIXED_FACTS — silently forcing decoder = "one_shot",
+# particle_type.target = "physical" and the v0.2 default sizes, while still
+# passing validate_config.
+config_version = 3
+
+[conditioning]
+# Physical-property MLPs rather than learned vocab embeddings: computable for
+# any PDG code / material, which is what the held-out-species and
+# held-out-material generalization comparisons need.
+out_dim      = 128
+share_stages = false
+
+# n_layers = 2 rather than the v0.3 default of 1: v0.2's conditioning MLP was
+# always 2 deep (see _migration.V02_FIXED_FACTS), so this keeps the encoder
+# identical to the architecture that produced the results cited above.
+[conditioning.particle]
+type     = "physical"
+emb_dim  = 16
+n_layers = 2
+
+[conditioning.material]
+type     = "physical"
+emb_dim  = 16
+n_layers = 2
+
+[stage1_model]
+generator    = "flow"
+hidden_dim   = 512
+n_res_blocks = 6
+dropout      = 0.0
+
+[stage2_model]
+# The v0.3 pivot: autoregressive in descending-energy order with a
+# categorical species target, which is the agreed response to the 2026-08-03
+# secondary-species failure. Flow (not the schema default wgan) so the
+# baseline varies only the decoder relative to the best v0.2 result.
+#
+# COST, measured (RTX 4070, bs 4096, 10 ODE steps), not estimated:
+# sample.sample_secondaries_ar loops `for k in range(k_max)` unconditionally
+# — all 15 slots regardless of predicted n_sec — so a flow AR token costs
+# k_max * steps = 150 stage-2 calls per physics step. That makes this block
+# the dominant cost on both sides:
+#   training   flow AR 29.5k samp/s  vs flow one-shot 190.7k samp/s  (6.5x)
+#   inference  flow AR  8.5k step/s  vs flow one-shot  68.7k step/s  (8.1x)
+# Accepted deliberately: one-shot is the configuration whose secondary
+# species distribution failed, and that failure is what v0.3 exists to fix.
+decoder      = "autoregressive"
+generator    = "flow"
+hidden_dim   = 512
+n_res_blocks = 6
+dropout      = 0.0
+k_max        = 15
+
+[stage2_model.autoregressive]
+history         = "markov"
+teacher_forcing = "always"
+
+[stage2_model.particle_type]
+target = "onehot"
+# Decoupled from conditioning.particle.emb_dim (gitea #29). 32 classes + the
+# "other" bucket keeps essentially all real secondary species out of "other"
+# without making the head expensive.
+n_classes    = 32
+other_policy = "sample"
+
+[train]
+epochs        = 50
+# Sized for ONE NVIDIA L40S on deepthought2 (46068 MiB; the box has two, and
+# CLAUDE.md's shared-machine rule allows a single GPU). From a measured
+# linear fit of this exact config's training step on the local RTX 4070:
+#   peak reserved MiB = 0.9736 * batch_size + 115
+# so 36864 reserves ~36.0 GiB, i.e. 78% of the card, leaving ~10 GiB of
+# headroom for fragmentation and the CUDA context. Throughput is already
+# flat above bs~4096 on the 4070, so this is chosen for occupancy on the
+# larger card, not for step efficiency — and it sits next to the 43008/32768
+# of the runs lr = 3e-4 was proven at.
+batch_size    = 36864
+lr            = 3e-4
+warmup_epochs = 3
+weight_decay  = 0.01
+ema_decay     = 0.9999
+val_fraction  = 0.1
+num_workers   = 4
+seed          = 0
+# The marginal/KL pass is expensive (~5000 s on top of an epoch), so keep it
+# to every 10th epoch; the cheap per-epoch val loss still runs every epoch.
+validate_every = 10
+validate_steps = 10
+wandb          = true
+wandb_project  = "giant"
@@ -0,0 +1,68 @@
+# Example GIANT workflow spec — `giant workflow run configs/workflow_example.toml`.
+#
+# One file parameterises a whole experiment: the datasets, the geometry oracle,
+# N trainings, N rollouts, and the analyses comparing them. Every task's output
+# directory carries a hash of its resolved sub-spec (plus its parents), so
+# editing anything here re-runs exactly the affected subtree and nothing else.
+#
+# result_dir/log_dir must be visible from both the submit host and the workers
+# (i.e. on /ceph) — there is deliberately no HTCondor file transfer.
+
+[workflow]
+name       = "baseline-vs-router"
+result_dir = "/ceph/lbogner/workflows/baseline-vs-router"
+log_dir    = "/ceph/lbogner/workflows/baseline-vs-router/logs"
+
+[condor]
+accounting_group = "cms"
+repo_dir         = "/work/lbogner/giant"   # also b2luigi's working_dir
+env_script       = "/work/lbogner/giant/condor_env.sh"
+docker_image_cpu = "cverstege/alma9-gridjob"
+docker_image_gpu = "mschnepf/slc7-condocker"
+remote           = true
+
+[dataset]
+steps     = "/ceph/lbogner/geant_steps/train/"     # training data
+reference = "/ceph/lbogner/geant_steps/holdout/"   # rollout seeds + analysis truth
+
+[geometry]
+method    = "slab"
+subsample = 500_000
+
+[[train]]
+name           = "baseline"
+config         = "configs/baseline.toml"
+epochs         = 200
+epochs_per_job = 1            # one short GPU job per epoch, chained
+request_gpus   = 1
+gpu_memory_mb  = 20000
+overrides      = { lr = 3e-4 }   # `giant train` flag names
+
+[[train]]
+name           = "router-balanced"
+config         = "configs/router.toml"
+epochs         = 200
+epochs_per_job = 1
+request_gpus   = 1
+gpu_memory_mb  = 20000
+
+[[rollout]]
+name          = "baseline"
+train         = "baseline"     # -> [[train]].name
+n_events      = 2000
+energy_cutoff = 0.1
+
+[[rollout]]
+name          = "router-balanced"
+train         = "router-balanced"
+n_events      = 2000
+energy_cutoff = 0.1
+
+[[analysis]]
+name        = "baseline-vs-router"
+rollouts    = ["baseline", "router-balanced"]
+chunks      = 32
+energy_bins = 4
+bins        = 50
+top_pdg     = 6
+gallery     = true
@@ -1,995 +0,0 @@
-# GIANT v0.3.0 — Stage-2 autoregressive redesign
-
-**Status:** design agreed, not implemented. Branch `v0.3.0-stage2-autoregressive`.
-**Date:** 2026-08-04.
-**Source:** `~/knowledge-base/meetings/2026-08-04-jan-stage2-autoregressive-architecture.md`
-(meeting with Jan), plus the design decisions taken in the session that produced
-this document.
-
-This document is the implementation contract for v0.3.0. It specifies the new
-config format option by option, the `network.py` refactor, and the order in which
-to build it. Read it before touching `giant/config.py` or `giant/model/network.py`.
-
---
-
-## Table of contents
-
-1. [Why](#1-why)
-2. [Decisions register](#2-decisions-register)
-3. [Config format reference](#3-config-format-reference)
-4. [Migration: v0.2 -> v0.3](#4-migration-v02---v03)
-5. [network.py refactor](#5-networkpy-refactor)
-6. [Stage-2 autoregressive design](#6-stage-2-autoregressive-design)
-7. [Training loop](#7-training-loop)
-8. [Data and setup-cache changes](#8-data-and-setup-cache-changes)
-9. [Config machinery changes](#9-config-machinery-changes)
-10. [Callers that need updating](#10-callers-that-need-updating)
-11. [Settled scope and open questions](#11-settled-scope-and-open-questions)
-12. [Implementation order](#12-implementation-order)
-
---
-
-## 1. Why
-
-The 2026-08-03 WGAN rollout benchmark
-(`~/knowledge-base/experiments/giant-wgan-physical-rollout-validation.md`) was
-good at the primary-step level and **failed at the secondary-species level**:
-zero photon secondaries generated, ~4M hallucinated `-14` (muon antineutrino)
-secondaries — a species essentially absent from Geant4.
-
-Stage 1 is not implicated; the meeting scoped everything to Stage 2. Two changes
-were agreed together:
-
- **Autoregressive generation** over secondaries in decreasing energy order,
-  replacing the one-shot masked `K_MAX=15` prediction, trained with teacher
-  forcing.
- **Categorical particle type** with a data-derived "other" bucket, reversing the
-  2026-07-17 move to a continuous `(log-mass, charge)` target. Working hypothesis:
-  the continuous target is part of what let the generator collapse onto degenerate
-  species.
-
-The meeting also set a **methodology**: compare Stage-2 architectures *standalone*
-(trained directly on secondary columns, no Stage-1 forward pass) before chaining
-the winner behind Stage 1. Iterating on the compounding-rollout-error problem is
-far cheaper that way than paying for a full two-stage run per candidate.
-
-That methodology is what forces the config refactor: v0.2 has a single global
-`train.mode` and a single `[model]` block, with no way to express "train stage 2
-only", "stage 1 flow + stage 2 WGAN", or "one-shot vs autoregressive stage 2".
-
---
-
-## 2. Decisions register
-
-| # | Decision | Rationale |
-|---|----------|-----------|
-| 1 | **`n_sec` head moves from stage 1 to stage 2** | Stage 1 becomes the pure 9D primary step. A stage-2-only run is then self-contained (it can predict its own multiplicity), and §3's "implicit stop token" alternative gets a natural home in the same config block. |
-| 2 | **Full mixed per-stage objectives** | `stage1 = flow` + `stage2 = wgan` must actually run — Stage 1 is good as flow, Stage 2 is what is being iterated on. The train loop becomes one trainer object per stage, each owning its optimizers and update cadence. |
-| 3 | **Migration shim for configs *and* checkpoints** | Nothing on `/ceph` goes dead. v0.2 `config.toml` files and v0.2 `model_config` dicts are translated on load. |
-| 4 | **Nested objective sub-tables** | `[stage1_model.wgan]`, `[stage2_model.ddpm]` rather than flat `wgan_n_critic` keys — self-documenting about which keys the active generator ignores, and validation can warn on a populated sub-table that is never read. |
-| 5 | **Particle type is adversarial: straight-through Gumbel into the critic** | The critic sees a relaxed one-hot alongside energy/direction, so the joint (species, kinematics) distribution is learned rather than factorized. Makes the collapse hypothesis directly testable instead of assumed. Accepts a **biased** gradient through the type path — assumed negligible, with a validation obligation in §11.4. |
-| 6 | **No charge conservation in v0.3.0** | Explicitly "not yet worked out" in the meeting. No `[stage2_model.conservation]` block at all. Energy conservation stays exact and implicit in the stick-breaking encoding. |
-| 7 | **Explicit stage-prefixed CLI flags** | `--stage2-hidden-dim` etc., no generic `--set path=value`. Discoverable via `--help` and tab-completable; the cost is a flag list kept in sync with `DEFAULT_CONFIG` by hand. |
-| 8 | **AR history is a config axis, markov default** | Markov (previous token + remaining budget + slot index) is the baseline that makes the meeting's §6 "is attention useful" question answerable by ablation rather than by comparing differently-shaped models. Both sit behind one `history_encoder(prefix) -> vector` interface. |
-
-### 2.1 Consequence of decision 5
-
-Straight-through Gumbel needs a critic to receive the relaxed one-hot. So the
-type mechanism is **implied by the generator**, and needs no config key of its own:
-
-| `stage2_model.generator` | Type mechanism |
-|--------------------------|----------------|
-| `"wgan"` | The type slice goes into the critic's input alongside energy/direction. Adversarial; learns the joint. Under `particle_type.target = "onehot"` it is relaxed through ST-Gumbel first; `"physical"` and `"embedding"` are already continuous and feed the critic directly. |
-| `"flow"` / `"ddpm"` | No critic exists -> the type slice trains against its own target, weighted by `particle_type.lambda` (same pattern as `n_sec` today): cross-entropy for `"onehot"`, regression for `"physical"` / `"embedding"`. Non-adversarial. |
-
-There is therefore **no `particle_type.adversarial` key** — the mechanism follows
-from `generator` × `particle_type.target`.
-
-### 2.2 Rejected alternatives worth remembering
-
- **`tie_to_stage1` as a tri-state** (`none`/`gate`/`full`). Sharing experts
-  between stages is impossible — the stage-1 trunk's input is the 9D target
-  vector, stage 2's is a token vector of a different width. Sharing the *gate*
-  is the only meaningful tying, so the key is a bool.
- **Generic `--set path.to.key=value` CLI overrides.** Considered and rejected in
-  favour of explicit flags (decision 7).
- **Per-expert trunk sizing** (`expert_hidden_dim` / `expert_n_blocks`). Removed
-  in v0.3.0: experts always use the stage's own `hidden_dim` / `n_res_blocks`.
-  This was already the effective behaviour — v0.2's `0` sentinel meant "inherit"
-  and nothing ever set it otherwise. Consequence worth knowing: a routed model is
-  `n_experts` × the parameters of the monolith at **equal per-row eval cost**
-  (top-1 dispatch runs one full-size trunk), so routing buys specialization, not
-  a per-call speedup.
-
-### 2.3 Correction to an earlier claim
-
-An earlier draft of this design asserted that a one-hot conditioning mode is
-"mathematically identical to an `nn.Embedding` lookup" and should be dropped.
-**That is wrong for the mode specified in §3.1.** One-hot here is a *fixed,
-unlearned* vector of width `emb_dim` covering the top `emb_dim - 1` species by
-training-set count plus an "other" bin. The difference from `embedding` is the
-**vocabulary cap**, not the parameterization: with 237 PDG codes and
-`emb_dim = 16`, `embedding` gives 237 distinct learned vectors while `onehot`
-gives 16 classes. That is a real capacity difference and a real statement about
-how rare species are treated, so all three modes are kept.
-
---
-
-## 3. Config format reference
-
-Six top-level blocks: `[conditioning]`, `[stage1_model]`, `[stage2_model]`,
-`[train]`, plus per-stage sub-tables and the existing `[meta]` (written by
-`save_config`, never hand-authored).
-
-### 3.1 `[conditioning]`
-
-One block shared by both stages. Each stage still builds its own encoder
-*instance* (separate weights) unless `share_stages = true`.
-
-**The particle and material axes are configured independently and may mix
-freely** — e.g. material `physical` with particle `embedding` is a valid and
-intended combination.
-
-```toml
-[conditioning]
-out_dim      = 128
-share_stages = false
-
-[conditioning.particle]
-type     = "physical"
-emb_dim  = 16
-n_layers = 1
-
-[conditioning.material]
-type     = "physical"
-emb_dim  = 16
-n_layers = 1
-```
-
-| Key | Type | Default | Meaning |
-|-----|------|---------|---------|
-| `out_dim` | int | `128` | Width of the fused conditioning vector produced by the encoder's fusion MLP, consumed by every downstream trunk. **New in v0.3.0** — v0.2 hardcoded this as `cond_out_dim = 128` in every constructor signature, unreachable from config. |
-| `share_stages` | bool | `false` | `false`: stage 1 and stage 2 each construct their own `ConditionEncoder` with identical config but independent weights (v0.2 behaviour). `true`: one instance, shared by reference. Shared weights halve the conditioning parameter count and force a common representation; independent weights let each stage specialize its view of the pre-step state. |
-
-#### `[conditioning.particle]` and `[conditioning.material]`
-
-Identical key sets, applied to the two identity axes independently.
-
-| Key | Type | Default | Meaning |
-|-----|------|---------|---------|
-| `type` | `"physical"` \| `"embedding"` \| `"onehot"` | `"physical"` | How this axis's identity becomes an `emb_dim`-wide vector. See the table below. |
-| `emb_dim` | int | `16` | Width of this axis's vector. Under `"onehot"` it **also sets the class count** — see below. |
-| `n_layers` | int | `1` | Depth of the sub-MLP under `"physical"`. `1` is the single-layer net with `emb_dim` neurons. Ignored under `"embedding"` / `"onehot"`. |
-
-**The three modes:**
-
-| mode | particle input | material input | width | learned parameters |
-|------|----------------|----------------|-------|--------------------|
-| `"physical"` | `log(mass)`, `charge` (`PARTICLE_PHYS_DIM = 2`) | `Z_eff`, `A_eff`, `log(density)`, `log(X0)`, `log(λ_int)` (`MATERIAL_PHYS_DIM = 5`) | `emb_dim` | one `n_layers`-deep MLP with `emb_dim` neurons |
-| `"embedding"` | dense vocab index | dense vocab index | `emb_dim` | `nn.Embedding(vocab, emb_dim)` |
-| `"onehot"` | top `emb_dim - 1` PDG codes by training-set count, plus one "other" bin | top `emb_dim - 1` materials by count, plus "other" | `emb_dim` | **none** — a fixed vector |
-
- `"physical"` reads columns already present in `cond_cont[:, COND_DIM_BASE:]`
-  (see `giant.data.transforms.build_features`). It is computable for **any** PDG
-  code or material, which is what allows generalization beyond the training menu.
- `"embedding"` memorizes the training menu — the generalization-comparison
-  baseline, and the only mode that supports
-  `stage2_model.particle_type.target = "embedding"` (§3.3).
- `"onehot"` is a *fixed, unlearned* representation. It is **not** a
-  reparameterization of `"embedding"`: the difference is the vocabulary cap. With
-  237 PDG codes and `emb_dim = 16`, `"embedding"` gives 237 distinct learned
-  vectors while `"onehot"` gives 16 classes. Needs the same data-derived top-N
-  map as the stage-2 type target (§8).
-
-**Interaction:** `particle.type = "physical"` is incompatible with router types
-`"pdg"` and `"process"`, which build their own dataset-scoped
-`nn.Embedding(pdg_vocab, ...)` regardless of the trunk's conditioning mode.
-Pairing them silently reintroduces a training-menu-scoped lookup at the routing
-layer, defeating the point of physical conditioning. Rejected loudly at build time
-— `_check_router_conditioning_compat` in `network.py`, which carries over but must
-now read `conditioning.particle.type` rather than a single global mode.
-
-### 3.2 `[stage1_model]`
-
-Stage 1 predicts the 9D primary post-step vector
-(`giant/constants.py:LOCAL_TARGET_NAMES`). As of decision 1 it carries **no**
-`n_sec` head.
-
-```toml
-[stage1_model]
-active       = true
-generator    = "flow"
-hidden_dim   = 256
-n_res_blocks = 6
-dropout      = 0.0
-lambda       = 1.0
-```
-
-| Key | Type | Default | Meaning |
-|-----|------|---------|---------|
-| `active` | bool | `true` | `false` skips building and training stage 1 entirely. The resulting checkpoint holds only stage 2 and **cannot be rolled out** — `giant rollout` must refuse it with a clear error. Used for the meeting's Stage-2-only architecture comparison. |
-| `generator` | `"flow"` \| `"ddpm"` \| `"wgan"` | `"flow"` | The generative objective. `"flow"`: conditional flow matching (Lipman et al. 2022), ~10 ODE steps at inference. `"ddpm"`: cosine-schedule diffusion baseline. `"wgan"`: WGAN-GP, single forward pass at inference. Replaces the global `train.mode`. Objective-specific knobs live in the matching sub-table below. |
-| `hidden_dim` | int | `256` | Trunk width — also the width of **every expert** under a routed trunk. |
-| `n_res_blocks` | int | `6` | Number of `ResBlock`s in the trunk, and in every expert under a routed trunk. Was `model.n_blocks`; renamed for clarity since v0.3.0 also has attention layers in stage 2. |
-| `dropout` | float | `0.0` | Dropout inside each `ResBlock`. **Default changed in v0.3.0** (was `0.1`). |
-| `lambda` | float | `1.0` | Weight of this stage's loss in the total. Meaningful when both stages are active and non-adversarial; a WGAN stage's adversarial loss drives its own optimizer, so `lambda` scales only its non-adversarial auxiliary terms. |
-
-#### `[stage1_model.flow]` — read only when `generator = "flow"`
-
-| Key | Type | Default | Meaning |
-|-----|------|---------|---------|
-| `time_dim` | int | `64` | Width of the `SinusoidalEmbedding` for the flow time variable, concatenated into the trunk's conditioning. **New in v0.3.0** — v0.2 hardcoded 64. |
-
-#### `[stage1_model.ddpm]` — read only when `generator = "ddpm"`
-
-| Key | Type | Default | Meaning |
-|-----|------|---------|---------|
-| `time_dim` | int | `64` | As above, for the diffusion time variable. |
-| `n_steps` | int | `1000` | Cosine-schedule diffusion steps. **New in v0.3.0** — v0.2 hardcoded this in `CosineSchedule`. |
-
-#### `[stage1_model.wgan]` — read only when `generator = "wgan"`
-
-| Key | Type | Default | Meaning |
-|-----|------|---------|---------|
-| `noise_dim` | int | `64` | Width of the generator's input noise vector. There is no time variable, hence no `time_dim`. |
-| `n_critic` | int | `5` | Critic updates per generator update (Gulrajani et al. 2017). |
-| `gp_weight` | float | `10.0` | Gradient-penalty coefficient. |
-| `critic_lr` | float | `0.0` | Critic learning rate. `0.0` means "inherit `train.lr`" — not `None`, since the TOML writer has no null literal to round-trip. |
-| `critic_hidden_dim` | int | `0` | Critic trunk width. `0` = inherit `stage1_model.hidden_dim`. **New in v0.3.0** — v0.2 always sized the critic from the generator. |
-| `critic_n_res_blocks` | int | `0` | Critic trunk depth. `0` = inherit `stage1_model.n_res_blocks`. |
-
-#### `[stage1_model.router]`
-
-Content carries over from v0.2's `[model.router]` unchanged. Reproduced here in
-full because the block is now per-stage and its semantics are easy to lose.
-
-```toml
-[stage1_model.router]
-enabled = false
-type    = "energy"
-n_experts = 4
-temperature = 0.5
-learn_centers = true
-learn_width = false
-learn_temperature = false
-width_min_ratio = 0.1
-width_max_ratio = 10.0
-lambda_balance = 0.0
-lambda_entropy = 0.0
-lambda_proc    = 0.0
-gumbel = false
-gumbel_tau_start = 1.0
-gumbel_tau_end   = 0.1
-emb_dim = 8
-hidden_dim = 64
-```
-
-| Key | Type | Default | Meaning |
-|-----|------|---------|---------|
-| `enabled` | bool | `false` | Replace the monolithic trunk with a mixture of per-expert trunks: soft-mixed over all experts at train time, **top-1 dispatched at eval time** (each row runs exactly one expert). Every expert is `hidden_dim` × `n_res_blocks` — v0.3.0 removes the per-expert sizing keys, so a routed model costs `n_experts` × the monolith's parameters at equal per-row eval cost. Routing buys specialization, not a per-call speedup. |
-| `type` | str | `"energy"` | Router impl from `ROUTER_REGISTRY`. `"energy"`: soft turn-on gate over normalized pre-step log-energy. `"pdg"`: gate over a learned PDG embedding. `"process"`: own classifier over pre-step conditioning predicting the step-ending physics process. `"composed"`: joint outer-product gating over multiple axes via `axis{i}_{field}` keys. |
-| `n_experts` | int | `4` | Number of experts. For `type = "process"` this doubles as the number of process classes. Ignored for `type = "composed"` (each axis has its own). |
-| `temperature` | float | `0.5` | Softmax denominator for the distance-based gate. As `tau -> 0` the gate hardens to nearest-center (Voronoi) selection, which is exactly what eval-time `top1` uses. Energy/pdg routers only. |
-| `learn_centers` | bool | `true` | Whether gate centers are `nn.Parameter` or a fixed buffer. |
-| `learn_width` | bool | `false` | Give each expert its own learnable width, sigmoid-bounded to `[width_min_ratio, width_max_ratio] * temperature`. Mutually exclusive with `learn_temperature`. |
-| `learn_temperature` | bool | `false` | Make the single shared `temperature` learnable, same bounding. Mutually exclusive with `learn_width`. |
-| `width_min_ratio` | float | `0.1` | Lower bound multiplier for the above. Must bracket 1.0 with `width_max_ratio` so enabling either mode is a no-op at init. |
-| `width_max_ratio` | float | `10.0` | Upper bound multiplier. **Deliberately bounded rather than `softplus`/`exp`:** an unbounded width lets one expert's logit `-d²/width -> 0` almost everywhere, so it wins nearly every row regardless of distance — the same "experts overlap instead of partitioning" failure the router design exists to avoid. |
-| `lambda_balance` | float | `0.0` | Importance-CV² load-balancing auxiliary loss weight (Shazeer et al. 2017). **The 2026-07-22 benchmark failure ran with `0.0`; do not repeat that.** |
-| `lambda_entropy` | float | `0.0` | Entropy-regularization weight penalizing uniform/collapsed gating. Secondary guard against all experts' widths co-inflating together, which `lambda_balance` cannot see (usage shares stay even throughout that failure). Use with caution: indiscriminate entropy penalties also suppress legitimate ambiguity near a decision boundary. |
-| `lambda_proc` | float | `0.0` | Supervised process-classification CE weight. `"process"` router only; `0.0` still trains a working router (the gate gets gradient through the downstream loss) but only `> 0` grounds it in the true `process` label. |
-| `gumbel` | bool | `false` | Straight-through Gumbel-softmax train-time combine weights: the forward pass samples a hard one-hot combination (matching eval-time top-1 dispatch exactly) while the backward pass still flows smooth gradient to every expert. Targets the train/eval mismatch. |
-| `gumbel_tau_start` | float | `1.0` | Gumbel temperature at step 0, annealed linearly over training. |
-| `gumbel_tau_end` | float | `0.1` | Gumbel temperature at the final step. |
-| `emb_dim` | int | `8` | The router's **own** pdg (and material) embedding width, separate from the trunk's `ConditionEncoder`. `"pdg"` / `"process"` routers only. |
-| `hidden_dim` | int | `64` | The `"process"` router's internal classifier hidden width. |
-
-**Composed routers** use flat `axis{i}_{field}` keys instead of `type`/`n_experts`
-— e.g. `axis0_type = "energy"`, `axis0_n_experts = 4`, `axis1_type = "pdg"`,
-`axis1_n_experts = 3`, `axis1_emb_dim = 8`. Indices must be contiguous from 0.
-Flat keys keep the block a table of scalars, which the merge machinery relies on.
-
-**`centers_init`** is not authored by hand: `giant/pipeline.py` populates it for
-`type = "energy"` from real data quantiles collected during the existing
-normalizer-fitting pass, then writes it into the checkpoint's `model_config`.
-
-### 3.3 `[stage2_model]`
-
-Stage 2 predicts `n_sec` and the per-secondary energy/direction/type.
-
-```toml
-[stage2_model]
-active         = true
-decoder        = "autoregressive"
-generator      = "wgan"
-hidden_dim     = 256
-n_res_blocks   = 6
-dropout        = 0.0
-lambda         = 1.0
-k_max          = 15
-context_dim    = 64
-stage1_context = "truth"
-```
-
-| Key | Type | Default | Meaning |
-|-----|------|---------|---------|
-| `active` | bool | `true` | `false` trains stage 1 alone. The checkpoint then has no secondary decoder; `giant rollout` must refuse it, `giant predict` still works. |
-| `decoder` | `"one_shot"` \| `"autoregressive"` | `"autoregressive"` | `"one_shot"`: predict all `k_max` slots simultaneously with padded slots masked from the loss — v0.2 behaviour, kept as the baseline arm of the meeting's §7 comparison. `"autoregressive"`: emit one secondary at a time in descending-energy order. |
-| `generator` | `"flow"` \| `"ddpm"` \| `"wgan"` | `"wgan"` | As stage 1. Under `"autoregressive"` this is the objective for **each token**: a WGAN token costs one forward pass, a flow token costs ~10 ODE steps. See the cost note in §6.4. |
-| `hidden_dim` | int | `256` | Trunk width. |
-| `n_res_blocks` | int | `6` | Trunk depth. |
-| `dropout` | float | `0.0` | Dropout inside each `ResBlock`. **Default changed in v0.3.0** (was `0.1`). |
-| `lambda` | float | `1.0` | Weight of stage 2's loss in the total. Was `train.lambda_s2`. |
-| `k_max` | int | `15` | Maximum secondary slots. Under `"one_shot"` this is the fixed output width; under `"autoregressive"` it is a safety cap on the generation loop. Was the global constant `K_MAX` in `giant/constants.py` (max observed `n_sec` is 14 in the PbWO4 dataset, so 15 covers it with one spare). |
-| `context_dim` | int | `64` | Width of the projected stage-1 outcome fed into stage 2's conditioning. Was the hardcoded `stage1_proj_dim = 64`. |
-| `stage1_context` | `"truth"` \| `"sampled"` | `"truth"` | What stage 2 conditions on during training. `"truth"`: the ground-truth stage-1 target vector, detached — v0.2 behaviour (`train.py:254` passes `x1_s1.detach()`), i.e. stage-level teacher forcing. `"sampled"`: stage 1's own sampled output, closing the train/inference gap at the cost of a sampling pass per batch and a moving target early in training. |
-
-#### `[stage2_model.n_sec]`
-
-| Key | Type | Default | Meaning |
-|-----|------|---------|---------|
-| `mode` | `"head"` \| `"stop_token"` \| `"truth"` | `"head"` | `"head"`: a classifier over `{0..k_max}` on the condition encoding alone (no diffusion noise), so it is callable independently at inference — v0.2 behaviour, and what the meeting's §3 explicitly decided to keep. `"stop_token"`: an EOS-style implicit stop — **accepted by the schema but not implemented in v0.3.0**, raising a clear "not implemented" error if set (§11.2). The key exists now so landing the mechanism later is not a config break. `"truth"`: take `n_sec` from ground truth — only valid for standalone stage-2 evaluation, never for rollout. |
-| `lambda` | float | `0.1` | Cross-entropy weight for the head. Was `train.lambda_nsec`. |
-
-#### `[stage2_model.particle_type]`
-
-**The three targets mirror the three conditioning modes of §3.1**, and use the
-same names.
-
-| Key | Type | Default | Meaning |
-|-----|------|---------|---------|
-| `target` | `"onehot"` \| `"physical"` \| `"embedding"` | `"onehot"` | What the token's type slice *is*. See the table below. |
-| `lambda` | float | `1.0` | Loss weight. Under `generator = "flow"`/`"ddpm"` this weights the cross-entropy (`"onehot"`) or regression (`"physical"`/`"embedding"`) term; under `"wgan"` the type is adversarial (§2.1) and this weights only any auxiliary term. |
-| `other_policy` | `"sample"` \| `"modal"` \| `"drop"` | `"sample"` | How a predicted "other" class becomes a concrete PDG code at rollout, needed because a secondary's mass/charge feed its own downstream conditioning. `"sample"`: draw from the empirical within-bucket distribution recorded at map-build time. `"modal"`: always the most common member. `"drop"`: discard the secondary. Read only under `target = "onehot"`. **Not decided in the meeting** — see §11. |
-
-**There is no `n_classes` key.** The class count under `target = "onehot"` is
-`conditioning.particle.emb_dim` — the same number that sizes the particle axis
-everywhere else. One knob sets the model's particle-type resolution, and the
-stage-2 onehot classes are by construction the same classes the conditioning
-onehot uses, so an emitted secondary's type is directly consumable as the
-conditioning of its own next step with no re-mapping.
-
-Note the coupling this creates: under `conditioning.particle.type = "physical"`,
-`emb_dim` primarily means "sub-MLP output width", yet it still sets the stage-2
-class count. Intentional, but worth knowing when tuning either.
-
-| target | token type slice | width | training target | inverse map at rollout |
-|--------|------------------|-------|-----------------|------------------------|
-| `"physical"` | regressed `(log mass, charge)` | `2` | the true PDG's physics values (`giant.particles.particle_mass_charge`) | none needed — mass/charge are used directly; `giant.particles.nearest_known_pdg` gives a reporting-only label. v0.2 behaviour. |
-| `"onehot"` | class logits | `conditioning.particle.emb_dim` | true class index | `argmax` -> class -> PDG (via `other_policy` for the "other" bin) |
-| `"embedding"` | an `emb_dim`-wide vector | `conditioning.particle.emb_dim` | `emb.weight[class].detach()` | L1-nearest row of `emb.weight` — see below |
-
-So the type slice is `conditioning.particle.emb_dim` wide for **both** `"onehot"`
-and `"embedding"`, and `2` only for `"physical"`.
-
-#### `target = "embedding"` in detail
-
-Stage 2 emits a vector that should equal **the conditioning's own particle
-embedding** for the secondary's species — the same `nn.Embedding` table
-`[conditioning.particle]` builds, not a second one.
-
-**Requires `conditioning.particle.type = "embedding"`.** There is no table to
-match against under `"physical"` or `"onehot"`; reject at config-validation time
-with an explicit error.
-
-**Why detached:** the regression target is `emb.weight[class].detach()`, so the
-embedding table receives gradient **only through the conditioning path**, never
-through the stage-2 output loss. Without the detach the target moves as the
-decoder chases it — the exact failure mode that motivated abandoning the learned
-type target in the first place (see `decisions/physical-property-conditioning`).
-The detach is what makes this option viable again.
-
-**Inverse map.** The natural exact-match form
-
-```python
-((out - emb.weight).abs().sum(1) < 1e-6).nonzero()
-```
-
-is correct as a **round-trip assertion in tests** (encode a known PDG, decode,
-recover the same PDG) but **must not be used at inference**: a generative model's
-continuous output essentially never lands within `1e-6` of a table row, so it
-returns an empty tensor almost always. Inference needs the nearest row:
-
-```python
-pdg_idx = (out.unsqueeze(-2) - emb.weight).abs().sum(-1).argmin(-1)   # L1 nearest
-```
-
-Note this decode is **unbounded in vocabulary**, unlike `"onehot"` — every PDG
-code in the training vocab is reachable, and there is no "other" bucket, hence no
-`other_policy`. The trade-off is that nearest-neighbour decode has no notion of
-confidence: an output far from every row still snaps to something.
-
-#### `[stage2_model.autoregressive]` — read only when `decoder = "autoregressive"`
-
-| Key | Type | Default | Meaning |
-|-----|------|---------|---------|
-| `order` | `"energy_desc"` | `"energy_desc"` | Canonical generation order. Descending energy is the ordering already flagged as natural in the Phase-2 note's open questions, and the one the existing stick-breaking encoding assumes. Single-valued for now; the key exists so an alternative ordering is not a config break. |
-| `history` | `"markov"` \| `"attention"` | `"markov"` | How token *i+1* sees tokens ≤ *i*. `"markov"`: previous token plus running scalars (remaining energy budget, slot index) — a fixed-width summary. `"attention"`: causal self-attention over all emitted tokens. See §6.2 for the trade-off. |
-| `teacher_forcing` | `"always"` \| `"scheduled"` \| `"never"` | `"always"` | `"always"`: condition on the ground-truth previous secondary throughout training (the meeting's confirmed plan). `"scheduled"`: scheduled sampling — interpolate toward conditioning on the model's own prediction. `"never"`: free-running from the start. |
-| `tf_p_start` | float | `1.0` | Under `"scheduled"`, P(use ground truth) at epoch 0. |
-| `tf_p_end` | float | `1.0` | Under `"scheduled"`, P(use ground truth) at the final epoch. Linear interpolation between the two. |
-| `attn_n_heads` | int | `4` | Attention heads. Read only under `history = "attention"`. |
-| `attn_n_layers` | int | `2` | Causal self-attention layers. Read only under `history = "attention"`. |
-
-#### `[stage2_model.flow]` / `[stage2_model.ddpm]`
-
-Same keys as their `[stage1_model.*]` counterparts (`time_dim`; plus `n_steps`
-for ddpm).
-
-#### `[stage2_model.wgan]` — read only when `generator = "wgan"`
-
-Same keys as `[stage1_model.wgan]` (`noise_dim`, `n_critic`, `gp_weight`,
-`critic_lr`, `critic_hidden_dim`, `critic_n_res_blocks`), plus:
-
-| Key | Type | Default | Meaning |
-|-----|------|---------|---------|
-| `gumbel_tau_start` | float | `1.0` | Straight-through Gumbel temperature for the **particle-type one-hot** at step 0, annealed linearly. Read only under `particle_type.target = "onehot"` (the other two targets are continuous and need no relaxation). Distinct from `router.gumbel_tau_start`, which anneals expert-combination weights — two unrelated Gumbel relaxations that must not share a key. |
-| `gumbel_tau_end` | float | `0.1` | Same, at the final step. |
-
-Under `"autoregressive"` a fresh `noise_dim` draw is made **per token**.
-
-#### `[stage2_model.router]`
-
-| Key | Type | Default | Meaning |
-|-----|------|---------|---------|
-| `tie_to_stage1` | bool | `false` | `true`: stage 2 shares **stage 1's `Router` module instance**, so expert *i* in stage 1 and expert *i* in stage 2 gate on identical conditions by construction. Every other key in this block is then ignored. `false`: an independent router — note this is v0.2's actual behaviour, which built *two separate routers from one config*, so stage-1 expert *i* and stage-2 expert *i* had no semantic relationship despite identical hyperparameters. |
-
-All other keys are as `[stage1_model.router]`. Invalid when `stage1_model.active
-= false` and `tie_to_stage1 = true` — reject at config-validation time.
-
-### 3.4 `[train]`
-
-Optimizer, schedule, data split and logging only. Everything model-shaped moved
-into the stage blocks.
-
-```toml
-[train]
-epochs = 100
-batch_size = 4096
-lr = 3e-4
-weight_decay = 0.01
-ema_decay = 0.9999
-warmup_epochs = 5
-val_fraction = 0.1
-max_val_batches = 200
-num_workers = 4
-seed = 0
-validate_every = 10
-validate_steps = 10
-wandb = true
-wandb_project = "giant"
-wandb_run_name = ""
-wandb_log_every = 50
-```
-
-| Key | Type | Default | Meaning |
-|-----|------|---------|---------|
-| `epochs` | int | `100` | Training epochs. |
-| `batch_size` | int | `4096` | Steps per batch. `auto` on the CLI estimates from free VRAM. |
-| `lr` | float | `3e-4` | AdamW learning rate for every stage's generator. |
-| `weight_decay` | float | `0.01` | AdamW weight decay. |
-| `ema_decay` | float | `0.9999` | EMA of model weights used for sampling; `0` disables. Maintained per stage. |
-| `warmup_epochs` | int | `5` | Linear LR warmup. |
-| `val_fraction` | float | `0.1` | Fraction of **events** (not steps) held out — the split is by `event_id` to avoid leaking correlated steps from the same shower. |
-| `max_val_batches` | int | `200` | Cap on the per-epoch val-loss pass; `0` = full val set. Distinct from `validate_every`'s marginal/KL pass. |
-| `num_workers` | int | `4` | DataLoader workers. Warns above ~1/4 of the machine's CPUs — portal machines are shared. |
-| `seed` | int | `0` | Seeds Python/numpy/torch and the event split. |
-| `validate_every` | int | `10` | Epochs between full marginal/KL validation passes. |
-| `validate_steps` | int | `10` | Sampler steps used during those passes. |
-| `wandb` | bool | `true` | W&B per-epoch metric logging. **Default changed in v0.3.0** (was opt-in `false`): the v0.3.0 work is a sequence of architecture comparisons, and a run that was not logged is not comparable. Set `false` for throwaway/debug runs. |
-| `wandb_project` | str | `"giant"` | W&B project. |
-| `wandb_run_name` | str | `""` | `""` means "use the checkpoint out_dir name" — not `None`, since the TOML writer has no null literal. |
-| `wandb_log_every` | int | `50` | Optimizer steps between batch-granularity metric logs. A single epoch can be tens of thousands of steps; per-epoch metrics always log in full. |
-
-### 3.5 Removed from v0.2
-
-| Key | Fate |
-|-----|------|
-| `train.mode` | Split into `stage1_model.generator` / `stage2_model.generator`. |
-| `train.lambda_nsec` | -> `stage2_model.n_sec.lambda`. |
-| `train.lambda_s2` | -> `stage2_model.lambda`. |
-| `train.n_critic`, `gp_weight`, `critic_lr` | -> `stage{1,2}_model.wgan.*`. |
-| `[model]` (whole block) | Split across `[conditioning]` and the two stage blocks. |
-
---
-
-## 4. Migration: v0.2 -> v0.3
-
-`[meta] config_version = 3` tags the new format; **absent means v0.2**.
-`migrate_config(cfg) -> cfg` applies the table below and is called on both
-`config.toml` load and checkpoint `model_config` load, so nothing on `/ceph` goes
-dead (decision 3).
-
-| v0.2 | v0.3 |
-|------|------|
-| `train.mode` | `stage1_model.generator` **and** `stage2_model.generator` (same value) |
-| `train.lambda_nsec` | `stage2_model.n_sec.lambda` |
-| `train.lambda_s2` | `stage2_model.lambda` |
-| `train.n_critic` / `gp_weight` / `critic_lr` | `stage1_model.wgan.*` **and** `stage2_model.wgan.*` |
-| `model.hidden_dim` / `n_blocks` / `dropout` | `stage{1,2}_model.hidden_dim` / `n_res_blocks` / `dropout` |
-| `model.emb_dim` | `conditioning.particle.emb_dim` **and** `conditioning.material.emb_dim` (v0.2 had one shared value) |
-| `model.conditioning` | `conditioning.particle.type` **and** `conditioning.material.type` (v0.2 had one shared mode) |
-| `model.noise_dim` | `stage{1,2}_model.wgan.noise_dim` |
-| `model.router.*` | `stage1_model.router.*`, copied verbatim to `stage2_model.router.*` with `tie_to_stage1 = false` (preserves v0.2's two-independent-routers behaviour) |
-| `model.expert_hidden_dim` / `expert_n_blocks` | **dropped.** v0.2's `0` sentinel meant "inherit from the monolith", which is now unconditional. A v0.2 config with a *non-zero* value must fail loudly rather than silently resize the experts — see §4.3. |
-| `model.k_max` | `stage2_model.k_max` |
-| — | `conditioning.out_dim = 128` (v0.2's hardcoded value) |
-| — | `conditioning.{particle,material}.n_layers = 2` (v0.2's hardcoded depth; note the v0.3 *default* is 1) |
-| — | `stage{1,2}_model.flow.time_dim = 64` / `.ddpm.time_dim = 64` (v0.2's hardcoded value) |
-| — | `stage2_model.context_dim = 64` (v0.2's hardcoded `stage1_proj_dim`) |
-| — | `stage2_model.decoder = "one_shot"` (v0.2 had no other option) |
-| — | `stage2_model.particle_type.target = "physical"` (v0.2 behaviour) |
-| — | `stage{1,2}_model.active = true` |
-
-### 4.1 The awkward one: `n_sec` head ownership
-
-Decision 1 moves the `n_sec` head to stage 2, but **v0.2 checkpoints carry
-`n_sec_head` weights inside the stage-1 module** (`DenoisingMLP.n_sec_head`,
-`WGANGenerator.n_sec_head`, `RoutedDenoisingMLP.n_sec_head`). The shim must keep
-those loading where they are.
-
-Handling: `migrate_config` sets an internal
-`stage2_model.n_sec.legacy_owner = "stage1"` that `build_models` honours by
-attaching the head to the stage-1 module. Never written by new runs, never
-CLI-settable, never documented as a user-facing option.
-
-### 4.2 Non-inheriting expert dims
-
-v0.3.0 drops `expert_hidden_dim` / `expert_n_blocks` (§2.2). Migration must
-distinguish two cases:
-
- value is `0` (the "inherit" sentinel, and what every real run used) — drop the
-  key silently, behaviour is unchanged.
- value is non-zero — **fail loudly.** Silently resizing those experts to
-  `hidden_dim` would change the architecture, so the checkpoint's weights would no
-  longer match. Such a checkpoint can only be loaded by v0.2.
-
-### 4.3 Migration test
-
-The acceptance criterion for the whole shim: **load a v0.2 checkpoint through
-`migrate_config` + the new `build_models`, and diff its outputs against v0.2 code
-on the same input batch.** Bit-identical, or the refactor has changed something it
-should not have. Pick one flow checkpoint and one WGAN checkpoint from `/ceph`.
-
---
-
-## 5. network.py refactor
-
-### 5.1 What it looks like today
-
-Ten classes that are permutations of three independent choices:
-
-| | flow/ddpm | wgan generator | wgan critic |
-|---|---|---|---|
-| **stage 1** | `DenoisingMLP` | `WGANGenerator` | `Critic` |
-| **stage 1, routed** | `RoutedDenoisingMLP` | — | — |
-| **stage 2** | `SecondaryDecoder` | `WGANSecondaryGenerator` | `SecondaryCritic` |
-| **stage 2, routed** | `RoutedSecondaryDecoder` | — | — |
-
-The empty cells are the entire reason `giant/pipeline.py:275` hard-rejects
-`--mode wgan --router`: no routed WGAN generator class was ever written. There is
-no deeper reason — the routed trunk is orthogonal to the objective.
-
-Every one of those classes repeats the same body: build a condition encoder,
-optionally a time embedding, project input, run blocks, project output.
-
-### 5.2 Proposed decomposition — one axis per config block
-
-**(a) `[conditioning]` -> encoders**
-
-```python
-ConditionEncoder(type, emb_dim, n_layers, out_dim)   # behaviour unchanged, now configurable
-ContextAdapter(in_dim, context_dim)                  # stage-1 outcome -> context vector
-```
-
-`SecondaryConditionEncoder` **disappears as a class**. It was
-`ConditionEncoder` + a `stage1_proj` linear + a fuse layer; those compose at the
-stage-model level instead.
-
-**(b) `[stage*_model]` + `.router` -> trunks, behind one interface**
-
-```python
-class Trunk(nn.Module):
-    def forward(self, x, cond, cond_cont=None, cond_cat=None) -> Tensor: ...
-
-MonolithicTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
-RoutedTrunk(router, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
-
-build_trunk(stage_cfg, in_dim, out_dim, cond_dim) -> Trunk
-```
-
-`ExpertTrunk` and `_route_forward` carry over unchanged. **One required change:**
-`MonolithicTrunk` and `ExpertTrunk` need a separate `out_dim` — today
-`ExpertTrunk` hardcodes `out_proj = nn.Linear(hidden_dim, in_dim)`, i.e.
-`out_dim == in_dim`. That stops working the moment stage 2's per-token output is
-`4 + type_dim` wide while its input is a noise vector of width `noise_dim`.
-
-**(c) `[stage*_model].generator` -> a thin wrapper, not a class family**
-
-The generator choice controls exactly two things:
-
- whether `SinusoidalEmbedding(t)` is concatenated into `cond` (flow/ddpm) or not (wgan)
- whether the trunk's `x` is the diffused/interpolated `x_t` (flow/ddpm) or a noise draw `z` (wgan)
-
-That is small enough for one wrapper, collapsing six of today's ten classes:
-
-```python
-class GenerativeTrunk(nn.Module):
-    """cond-encode -> (optional time-embed) -> trunk. Backs flow, ddpm and wgan."""
-    def forward(self, x, cond_cont, cond_cat, t=None, context=None) -> Tensor: ...
-```
-
-### 5.3 Resulting class list
-
-```
-# building blocks
-SinusoidalEmbedding, ResBlock, ConditionEncoder, ContextAdapter
-
-# trunks
-Trunk (ABC), MonolithicTrunk, RoutedTrunk, ExpertTrunk
-
-# routers — carried over unchanged
-Router, EnergyRouter, PdgRouter, ProcessRouter, ComposedRouter
-ROUTER_REGISTRY, register_router, build_router, build_composed_router
-
-# history encoders — new, stage-2 AR only
-HistoryEncoder (ABC), MarkovHistory, AttentionHistory
-
-# stage models
-Stage1Model            # 9D primary step
-Stage2OneShot          # k_max slots at once (v0.2 behaviour)
-Stage2Autoregressive   # one token at a time
-CriticModel            # stage-1 or stage-2 critic, generator-agnostic
-
-# factories
-build_models(cfg)  -> {"stage1": ... | None, "stage2": ... | None}
-build_critics(cfg) -> {"stage1": ... | None, "stage2": ... | None}
-```
-
-Ten classes become four stage classes plus reusable parts, and **routed WGAN comes
-for free** — `pipeline.py`'s rejection can be deleted.
-
-### 5.4 Factory signature change
-
-`build_models` returns a **dict, not a tuple**: `active = false` on either stage
-means that key is `None`. Every caller unpacking
-`stage1, sec_decoder = build_models(...)` must be updated
-(`pipeline.py:390`, `cli.py:871`, `cli.py:1254`).
-
---
-
-## 6. Stage-2 autoregressive design
-
-### 6.1 Token layout
-
-Secondaries are emitted one at a time in descending energy order. Per-token output
-width is `4 + type_dim`:
-
-| slice | meaning |
-|-------|---------|
-| `[0]` | stick-breaking logit — fraction of the **remaining** energy budget |
-| `[1:4]` | local-frame direction, normalized to a unit vector |
-| `[4:4+type_dim]` | particle type — see below |
-
-`type_dim` follows `particle_type.target` (§3.3): `2` for `"physical"`, and
-`conditioning.particle.emb_dim` for both `"onehot"` and `"embedding"`.
-Only `"onehot"` involves a relaxation — ST-Gumbel under `wgan`, cross-entropy
-under `flow`/`ddpm`; the other two are continuous and feed the critic (or the
-regression loss) directly.
-
-Per-token conditioning is:
-
-```
-base condition encoding          (ConditionEncoder, [conditioning].out_dim)
-+ stage-1 context                (ContextAdapter, stage2_model.context_dim)
-+ history_encoder(prefix)        (§6.2)
-+ running scalars                (remaining energy budget, slot index)
-```
-
-### 6.2 History encoders
-
-One interface, `history_encoder(prefix) -> fixed-width vector`:
-
- **`MarkovHistory`** — the previous token's `(energy_fraction, direction,
-  type_embedding)`. Fixed-width, one small MLP.
- **`AttentionHistory`** — causal self-attention over all emitted tokens, taking
-  the last position. `attn_n_heads` × `attn_n_layers`.
-
-**The trade-off** (this is decision 8's reasoning, recorded so it does not have to
-be re-derived):
-
-1. **Markov is less impoverished than it sounds.** The remaining-energy budget is
-   an *exact sufficient statistic* for the conservation constraint — stick-breaking
-   needs nothing else from history. Slot index likewise. What markov genuinely
-   cannot see is set *composition*: "I have already emitted two photons and an
-   electron." That matters for correlated production — pair production emits
-   exactly e⁺e⁻, a brems cascade correlates species across the set. Note that the
-   meeting's charge-conservation idea was exactly such a hand-engineered summary
-   statistic, and it is out of scope for v0.3.0 (decision 6), so markov does not
-   get that crutch.
-2. **Sequences are short and front-loaded.** `n_sec` is 0–2 for most steps, max 14.
-   For K ≤ 2, attention *is* markov. They diverge only in the high-multiplicity
-   tail — physically interesting but data-poor, so the attention path trains on
-   few examples where it actually matters.
-3. **Cost is not where intuition puts it.** Under teacher forcing both train in a
-   single parallel pass over all K tokens (every token's input is ground truth, so
-   nothing is sequential). At inference both need K sequential forwards; attention
-   additionally needs a KV cache to avoid re-encoding the prefix. Attention's
-   marginal FLOPs over ≤15 tokens are rounding error.
-4. **Exposure bias cuts against attention.** Attention conditions on the entire
-   generated prefix, so one off-manifold early token poisons every later token
-   through the context. Markov's fixed summary sees only one bad token, and the
-   budget scalars stay exact regardless. Given the explicitly-flagged train/
-   inference gap and the known compounding-rollout-error problem, the more
-   expressive history is also the more fragile one.
-
-Hence: markov is the default and the baseline; attention is a flag.
-
-### 6.3 Energy budget under AR
-
-**No re-derivation needed**, despite the meeting's action item suggesting
-otherwise. The existing stick-breaking is already sequential in spirit — each slot
-takes a fraction of what remains — so it carries over to per-token generation
-directly by feeding "remaining budget" as a per-token conditioning scalar.
-Conservation stays exact by construction: the valid slots' energies sum to `e_sec`,
-which the Stage-1 simplex already guarantees sums correctly with `edep` and
-`post_E`.
-
-### 6.4 Cost warning
-
-AR costs **K sequential forward passes per step** where one-shot costs one.
-Against the ~10× native-Geant4 eval budget that motivated the whole fast-eval
-track, this is the number to watch — not the history-encoder choice. With
-`generator = "flow"` it is worse still: ~10 ODE steps per token, so ~150 forwards
-per step in the worst case. `generator = "wgan"` (one pass per token) is the only
-configuration that plausibly meets the budget; flow AR is for quality comparison.
-
---
-
-## 7. Training loop
-
-Decision 2 (full mixed per-stage objectives) makes `train.py` one trainer object
-per active stage:
-
-```python
-class StageTrainer:                 # owns optimizers, EMA, update cadence
-    def step(self, batch, global_step) -> dict[str, float]: ...
-
-FlowTrainer, DDPMTrainer, WGANTrainer(critic, n_critic, gp_weight)
-```
-
- Non-adversarial stages contribute `lambda * loss` to one backward pass.
- A WGAN stage runs its own critic inner loop on the same batch, with a generator
-  update every `n_critic`-th batch — today's `_wgan_train_step` cadence.
- A mixed run (`flow` + `wgan`) steps stage 1 every batch while stage 2 does 5
-  critic updates then a generator update. Independent optimizers, independent EMA.
- Router auxiliary losses (`lambda_balance` / `lambda_proc` / `lambda_entropy`)
-  become per-stage, summed over whichever stages are routed. Today's
-  `hasattr(model, "router")` check (`train.py:262`) generalizes cleanly.
- `metrics.csv` and W&B metric names gain a stage prefix.
- With `stage1_model.active = false`, stage 2 still needs its stage-1 context: it
-  comes from the ground-truth target already in the batch (`x1_s1`), which is
-  exactly what v0.2 does anyway. **Stage-2-only training is therefore a cheap
-  ablation, not new plumbing** — drop the stage-1 loss, skip building stage 1.
-
---
-
-## 8. Data and setup-cache changes
-
-Three config options need a "top N−1 by training-set count plus other" map, but
-they resolve to **at most two distinct maps per run** — one per axis — because
-the class count always comes from that axis's `emb_dim`:
-
-| consumer | axis | N |
-|----------|------|---|
-| `conditioning.particle.type = "onehot"` | PDG | `conditioning.particle.emb_dim` |
-| `stage2_model.particle_type.target = "onehot"` | PDG | `conditioning.particle.emb_dim` |
-| `conditioning.material.type = "onehot"` | material | `conditioning.material.emb_dim` |
-
-The two PDG consumers therefore **share one map** — which is the point of
-dropping `n_classes`: a secondary's emitted type is directly consumable as the
-conditioning of its own next step, with no re-mapping between two class systems.
-Build one shared helper, structurally identical to today's `proc_map`:
-
- `build_topn_map_from_files(files, column, n_classes=...)` in
-  `giant/data/loader.py`, next to `build_process_map_from_files`
- a `setup_cache` section keyed by `(axis, N)`, same shape as `cache.proc_maps`
-  (which is keyed by `n_experts`). N is still part of the key so the sidecar stays
-  reusable across runs with different `emb_dim`, even though a single run only
-  ever needs one N per axis.
- persisted into the checkpoint beside `pdg_map` / `mat_map`
- inverted at rollout to recover a concrete PDG -> mass/charge per secondary
-
-Record the **empirical within-bucket distribution** at map-build time as well —
-`other_policy = "sample"` needs it.
-
-`particle_type.target = "embedding"` needs no map: it reaches the full training
-vocab through the conditioning's embedding table (§3.3).
-
-`giant/constants.py`: `K_MAX` and `SEC_DIM` stop being authoritative constants
-(they become `stage2_model.k_max` and a derived quantity). Keep them as defaults
-only, and audit the ~12 modules importing `K_MAX` for places that assume it is
-global truth.
-
---
-
-## 9. Config machinery changes
-
-All in `giant/config.py`:
-
- **`merge_cli_overrides`** — replace the hand-written one-level router merge with
-  a generic recursive deep-merge. The new layout is three levels deep
-  (`stage1_model.router.axis0_type`).
- **`save_config`** — recursive TOML writer; today it handles exactly one nesting
-  level (see its `nested_sections` list).
- **`default_out_dir_name`** — `_OUT_DIR_NAME_CANDIDATES` entries become dotted
-  paths (`"stage2_model.decoder"`) instead of `(section, field)` pairs. Add
-  candidates for the new discriminating fields: `decoder`, per-stage `generator`,
-  `particle_type.target`, `autoregressive.history`.
- **`resolve_expert_dims`** — **deleted.** Experts always take the stage's
-  `hidden_dim` / `n_res_blocks` (§2.2). Its two callers (`pipeline.py:351` and the
-  CLI's batch-size auto-estimate) read the stage keys directly, as does
-  `pipeline.py`'s "experts are NxM, different from model.hidden_dim" warning,
-  which becomes unreachable and should go.
- **`migrate_config(cfg) -> cfg`** — §4's table, applied on both `config.toml` load
-  and checkpoint `model_config` load. New `[meta] config_version = 3`.
- **`Conditioning` enum** — gains a third member `onehot`, and now feeds **two**
-  keys (`conditioning.particle.type`, `conditioning.material.type`) rather than
-  one. Shared with `scripts/dwarf.py`, so both CLIs stay in sync.
- **Cross-block validation** — a new `validate_config(cfg)` pass, since v0.3.0 has
-  constraints no single block can check:
-  `particle_type.target = "embedding"` requires
-  `conditioning.particle.type = "embedding"`; router types `"pdg"`/`"process"`
-  require `conditioning.particle.type != "physical"`;
-  `stage2_model.router.tie_to_stage1` requires `stage1_model.active`;
-  `n_sec.mode = "truth"` is invalid for a rollout-capable checkpoint.
- **`estimate_batch_size`** — its calibration constants assume the v0.2
-  architecture ("post-Phase-2, including the Stage-2 secondary decoder and n_sec
-  head"). AR stage 2 changes the activation-memory profile. **Recalibrate last**
-  (§12 step 8), measured on real hardware with the example configs.
-
---
-
-## 10. Callers that need updating
-
-| File | Why |
-|------|-----|
-| `giant/pipeline.py` | Builds `model_config`; now per-stage. Delete the wgan+router rejection at `:275`. Router `centers_init` seeding becomes per-stage. |
-| `giant/train.py` | Per-stage trainers (§7). |
-| `giant/cli.py` | Stage-prefixed flags for `train` and `new-run`; `build_models` now returns a dict (`:871`, `:1254`); `:1339` writes `model_config` into the rollout sidecar. |
-| `giant/sample.py` | Sampler picked per stage from `stage*_model.generator`; new AR sampling loop with KV cache under `history = "attention"`. |
-| `giant/rollout.py` | AR secondary generation; categorical class -> PDG decode; `other_policy` handling. |
-| `giant/validate.py` | Stage-2 marginals gain a type-class marginal. |
-| `giant/analysis/render.py` | `_router_summary(model_config)` at `:44` reads `model_config["router"]`. |
-| `giant/analysis/router_gating.py` | Same, at `:76`. |
-| `giant/constants.py` | `K_MAX` / `SEC_DIM` demoted to defaults (§8). |
-| `configs/*.toml` | All eight shipped configs are v0.2-format; regenerate or rely on the shim. |
-| `condor-gpu-train-rollout` branch | Submits `giant train` flags; needs rebasing onto the new flag surface. |
-
---
-
-## 11. Settled scope and open questions
-
-### 11.1 Settled — build as specified
-
- **`other_policy` switch.** The three-way `"sample"` / `"modal"` / `"drop"`
-  design in §3.3 is approved as the config surface for turning a predicted "other"
-  class into a concrete PDG at rollout. `"sample"` (draw from the empirical
-  within-bucket distribution recorded at map-build time) stays the default as the
-  least-biased option.
- **`[stage*_model.flow]` and `[stage*_model.ddpm]` stay separate tables.** The
-  duplicated `time_dim` is accepted; a merged `[stage*_model.diffusion]` was
-  considered and rejected because the name fits flow matching poorly.
-
-### 11.2 Deferred — not implemented in v0.3.0
-
- **`n_sec.mode = "stop_token"`.** The value is **accepted by the schema** but
-  raises a clear "not implemented in v0.3.0" error if set. The key existing now
-  means landing the implicit-stop mechanism later is not a config break. `"head"`
-  is what v0.3.0 builds, per the meeting's §3 decision.
- **Charge conservation.** No key at all — `[stage2_model.conservation]` is absent
-  from v0.3.0 (decision 6), unlike `stop_token` above. **Deliberately left
-  undesigned:** the mechanism should be worked out on its own terms when it is
-  taken up, not pre-shaped by choices made for this refactor. Adding the block
-  later is a config addition, not a break.
- **`stage2_model.generator = "ddpm"`.** The value is **accepted by the
-  schema** (§3.3 lists `"flow" | "ddpm" | "wgan"` with no caveat) but
-  `FlowDDPMStageTrainer.__init__` (`giant/train.py`) raises
-  `NotImplementedError` for stage 2 — only `"flow"` and `"wgan"` have a
-  stage-2 secondary-decoder loss implemented. `stage1_model.generator =
-  "ddpm"` is unaffected; this restriction is stage-2-only. Landing stage-2
-  ddpm later is a trainer addition, not a config break.
-
-### 11.3 Tracked as implementation work
-
- **Log the L1 distance distribution at rollout** under
-  `particle_type.target = "embedding"`. Nearest-neighbour decode has no reject
-  option — an output far from every table row still snaps to its nearest neighbour,
-  with no "other" bin and no confidence signal. A heavy tail in that distribution
-  means the decoder is emitting vectors off the embedding manifold, which is the
-  direct analogue of the species-collapse symptom this redesign exists to fix.
-  Belongs with the `giant/rollout.py` decode work (step 6) and should surface as a
-  `giant analyze` diagnostic plot alongside `router_gating`.
- **`estimate_batch_size` recalibration** for AR stage 2 is the **last** step of
-  the implementation flow (§12 step 8), measured on real hardware using the example
-  configs — not guessed from the existing calibration constants.
-
-### 11.4 Accepted with a validation obligation: differentiability
-
-**Position taken: differentiability through the categorical type path is broken,
-and that is accepted.** The expected contribution of the broken path to the total
-gradient is small enough to ignore. **This is an assumption, not a result — it has
-to be demonstrated later.**
-
-Recording it precisely, since "broken" covers three distinct things:
-
-| where | status |
-|-------|--------|
-| Per-token training loss under teacher forcing | **Fine.** Softmax cross-entropy needs no sampling; the loss is differentiable in the logits. |
-| ST-Gumbel into the critic (`generator = "wgan"`, §2.1) | **Biased, not absent.** The forward pass is a hard one-hot; the backward pass pretends it was the soft sample. Gradient flows, but it is not the gradient of what was actually computed. |
-| Full shower-rollout backprop | **Structurally broken regardless.** Already non-differentiable once secondaries spawn branches, independent of the type representation — so the categorical switch costs nothing that was not already lost. |
-
-The claim being accepted is about the middle row: the straight-through estimator's
-bias, propagated into the shared trunk, is expected to be negligible against the
-gradient from the continuous paths (stick-breaking energy, direction, and — when
-stage 1 is active — the 9D primary target). The third row is the reason this is
-tolerable at all: end-to-end differentiability was never available.
-
-**Validation obligation.** Do not treat this as settled until one of the following
-has actually been run. Cheapest first:
-
-1. **Gradient-magnitude accounting.** Instrument a training run to log the norm of
-   the trunk gradient contributed through the type slice against the norm from the
-   continuous slices. "Negligible" should mean a stable, small ratio — not merely
-   small at initialization. This is the direct measurement of the claim and costs
-   almost nothing to add.
-2. **Detached-type ablation.** Train with the type path detached from the shared
-   trunk entirely (type head still learns; no type gradient reaches the trunk)
-   against the ST-Gumbel default. Comparable species marginals and kinematics mean
-   the coupling was weak, which is the same conclusion by a different route.
-3. **Estimator swap**, only if 1–2 are inconclusive: compare ST-Gumbel against an
-   unbiased-but-high-variance estimator (e.g. REINFORCE with a baseline) on a short
-   run. Agreement in the learned marginals means the bias did not matter.
-
-Option 1 should be added when the AR trunk lands (step 5), so the evidence accrues
-during the architecture comparison rather than needing a dedicated run afterwards.
-
-**Why it still matters that this is written down:** decision 5 (adversarial type
-via ST-Gumbel) rests on this assumption. If the ratio in test 1 turns out not to be
-small, the fallback is not a redesign — it is `particle_type.target = "physical"`
-or a non-adversarial CE head, both of which already exist as config options.
-
---
-
-## 12. Implementation order
-
-1. **`config.py`** — new `DEFAULT_CONFIG`, recursive merge/write, `migrate_config`,
-   tests. Nothing else can land first.
-2. **`network.py`** — the §5 decomposition, with `Stage2OneShot` reproducing v0.2
-   exactly. Gate on the §4.3 migration test: load a v0.2 checkpoint through the
-   shim and diff outputs against v0.2 code.
-3. **`train.py`** — per-stage trainers; `active = false` paths. At this point
-   Stage-2-only training works and the meeting's step 2 (one-shot WGAN baseline,
-   trained standalone) is runnable.
-4. **Type map** — `loader.py` + `setup_cache.py` + checkpoint persistence +
-   `particle_type.target = "onehot"` in `Stage2OneShot`. This is the meeting's
-   action item 1, and it is testable against the one-shot baseline before any AR
-   work.
-5. **`Stage2Autoregressive`** with `history = "markov"`, `teacher_forcing =
-   "always"`. The meeting's step 3. Add the gradient-magnitude instrumentation of
-   §11.4 test 1 here, so the differentiability assumption accrues evidence during
-   the architecture comparison instead of needing its own run later.
-6. **`sample.py` / `rollout.py`** — AR generation and class -> PDG decode, so an AR
-   model can actually be rolled out and put through `giant analyze`. Includes the
-   L1-distance diagnostic of §11.3.
-7. **`history = "attention"`, scheduled sampling** — then run the meeting's §7
-   comparison (one-shot vs autoregressive, standalone) and only chain the winner
-   behind Stage 1.
-8. **`estimate_batch_size` recalibration** — last, once the architectures are
-   settled. Measure on real hardware with the example configs and record the new
-   calibration points in `config.py` (§9, §11.3).
-
-Steps 1–3 are pure refactor with a bit-identical acceptance criterion. Steps 4–7
-are the actual physics change. Step 8 is measurement, deliberately last: the
-activation-memory profile is not knowable until the AR trunk and its history
-encoder are final.
@@ -1,37 +0,0 @@
-# v0.3.0 — post-implementation audit: open discrepancies
-
-**Status:** steps 1–7 of `docs/v0.3.0-design.md` §12 are implemented (branch
-`v0.3.0-stage2-autoregressive`, commits `eb6dd27`..`200c6d2`). This document
-tracked discrepancies found between that implementation and the design contract
-during a 2026-08-07 audit, as concrete work items. **All items (1-9) are now
-resolved** — either implemented (1-5, 7-9) or explicitly deferred into
-`docs/v0.3.0-design.md` §11.2 (6: `stage2_model.generator = "ddpm"`). Step 8
-(`estimate_batch_size` recalibration) was intentionally still outstanding per
-§12 and was never tracked here.
-
---
-
-## Confirmed correct during the audit (no action needed)
-
-For reference — these were explicitly checked against the design doc and
-match it, including two spots the doc itself flagged as likely stale that
-turned out fine:
-
- Config schema, `DEFAULT_CONFIG`, `migrate_config` table (§4), `save_config`/
-  `merge_cli_overrides` recursion, `default_out_dir_name`, `Conditioning` enum,
-  `n_sec.mode = "stop_token"` error (§9, §11.2).
- `network.py`'s full class decomposition (§5.3), dict-returning
-  `build_models`/`build_critics` (§5.4), `ExpertTrunk` separate in/out dims,
-  ST-Gumbel wiring (§2.1), AR token layout (§6.1), Markov/Attention history
-  encoders (§6.2).
- Shared PDG top-N type map (one map, not two, per §8), `other_policy`
-  sample/modal/drop (§11.1), embedding L1-nearest decode, and the L1-distance
-  diagnostic surfaced in `giant analyze` (§11.3).
- `analysis/render.py`/`analysis/router_gating.py` correctly branch
-  old-flat vs new-nested `model_config["router"]` location — doc flagged this
-  as a likely stale spot (§10) but it's actually fine.
- `train.py`'s per-stage trainers, mixed flow+wgan runs, WGAN critic cadence,
-  per-stage router auxiliary losses, stage-prefixed metrics, stage-2-only
-  training via ground-truth `x1_s1` (§7).
- `pipeline.py`'s deleted wgan+router rejection, per-stage `centers_init`
-  seeding, removed stale expert-size warning (§9, §10).
@@ -0,0 +1,70 @@
+"""Shared v0.2 -> v0.3 migration knowledge.
+
+v0.3.0 broke the config format (single `[train]` + `[model]` -> `[conditioning]`/
+`[stage1_model]`/`[stage2_model]`/`[train]`), and that break has to be absorbed by two
+independent migration surfaces: `giant.config.migrate_config` (a v0.2 `config.toml`) and
+`giant.model.network._migrate_legacy_model_config` (a v0.2 checkpoint's flat
+`model_config` dict). Both translate the same v0.2 facts into the same v0.3 shape, so
+the facts live here once rather than as two hand-maintained copies — see issues.md
+Issue 6.
+
+A dependency-free leaf module so neither `config.py` nor `network.py` has to import the
+other to share this.
+"""
+
+# v0.2 model-shaped keys (config.toml's [model] table, or a checkpoint's flat
+# model_config dict — same key names in both) applied identically to both v0.3 stage
+# blocks, because v0.2 had only one trunk shape shared by both stages.
+V02_MODEL_KEY_TO_STAGES: tuple[tuple[str, str], ...] = (
+    ("hidden_dim", "hidden_dim"),
+    ("n_blocks", "n_res_blocks"),
+    ("dropout", "dropout"),
+)
+
+# v0.2 architectural facts that had no corresponding config key at all — always true of
+# a v0.2 model, so both migration surfaces inject them unconditionally. Keyed by dotted
+# path relative to the migrated dict's root. NOTE: conditioning.*.n_layers (2) differs
+# from the v0.3 *default* (1) — not a typo, v0.2's conditioning MLP was always 2 layers
+# deep.
+V02_FIXED_FACTS: dict[str, object] = {
+    "conditioning.out_dim": 128,
+    "conditioning.particle.n_layers": 2,
+    "conditioning.material.n_layers": 2,
+    "stage1_model.active": True,
+    "stage1_model.flow.time_dim": 64,
+    "stage1_model.ddpm.time_dim": 64,
+    "stage2_model.active": True,
+    "stage2_model.flow.time_dim": 64,
+    "stage2_model.ddpm.time_dim": 64,
+    "stage2_model.context_dim": 64,
+    "stage2_model.decoder": "one_shot",
+    "stage2_model.particle_type.target": "physical",
+}
+
+
+def reject_legacy_router_expert_sizing(router_cfg: dict, *, source: str) -> None:
+    """Pop and validate v0.2's per-expert width/depth override, in place.
+
+    v0.3.0 removed per-expert sizing — experts always inherit the stage's
+    hidden_dim/n_res_blocks — so a v0.2 router config/checkpoint that set a non-default
+    `expert_hidden_dim`/`expert_n_blocks` describes experts with a different width/depth
+    than the monolith, and can only be reproduced by v0.2 code. Silently dropping these
+    keys (a router builder's kwarg filtering would do this for free) would resize the
+    experts instead of refusing, so this raises loudly.
+
+    Always pops both keys, whether or not they were non-default, so callers can go on
+    to use the (now-cleaned) `router_cfg` unconditionally. `source` names what's being
+    migrated (e.g. "v0.2 config's model.router" or "this checkpoint's
+    model_config.router") for the error message.
+    """
+    expert_hidden_dim = router_cfg.pop("expert_hidden_dim", 0)
+    expert_n_blocks = router_cfg.pop("expert_n_blocks", 0)
+    if not (expert_hidden_dim or expert_n_blocks):
+        return
+    raise ValueError(
+        f"{source} sets expert_hidden_dim/expert_n_blocks to a non-default value "
+        f"({expert_hidden_dim!r}, {expert_n_blocks!r}); v0.3.0 removed per-expert "
+        "sizing (experts always inherit the stage's hidden_dim/n_res_blocks), so "
+        "this router's experts have a different width/depth than the monolith. "
+        "This checkpoint/config can only be loaded by v0.2 code."
+    )
@@ -1,9 +1,10 @@
 """Rollout-vs-reference analysis: streaming compute + plotstyle rendering.

-Compares one autoregressive ``giant rollout`` against a held-out miniCaloSim
-reference file, producing publication-styled comparison plots generated in
-parallel on HTCondor (one job per plot x data chunk, compute/merge/render
-split).
+Compares one or more autoregressive ``giant rollout`` runs against a single
+held-out miniCaloSim reference file shared by all of them, producing
+publication-styled comparison plots (one colored series per rollout, one
+reference line) generated in parallel on HTCondor (one job per plot x data
+chunk, compute/merge/render split) — orchestrated by ``giant/workflow``.

 Only ``render`` (and the ``render`` CLI path) imports plotstyle/LaTeX; everything
 re-exported here is plotstyle-free so it runs on a compute worker. Import
@@ -11,41 +12,42 @@ re-exported here is plotstyle-free so it runs on a compute worker. Import
 """

 from giant.analysis.catalog import build_catalog, catalog_ids, get_spec
-from giant.analysis.condor import (
+from giant.analysis.run import (
+    LoadedRollout,
    RunMeta,
-    SubmitConfig,
    compute_one,
    compute_reduced,
    derive_run_dir,
    load_rollout_yaml,
+    load_rollout_yamls,
    merge_all,
    merge_one,
    prep,
-    write_submit,
 )
 from giant.analysis.context import Context, build_context
 from giant.analysis.reduced import Partial, Reduced
 from giant.analysis.runtime_estimate import RUNTIME_SAFETY_MARGIN, estimate_runtime_s
-from giant.analysis.sources import Side
+from giant.analysis.sources import RolloutSpec, Side

 __all__ = [
    "build_catalog",
    "catalog_ids",
    "get_spec",
+    "LoadedRollout",
    "RunMeta",
-    "SubmitConfig",
    "compute_one",
    "compute_reduced",
    "derive_run_dir",
    "load_rollout_yaml",
+    "load_rollout_yamls",
    "merge_all",
    "merge_one",
    "prep",
-    "write_submit",
    "Context",
    "build_context",
    "Partial",
    "Reduced",
+    "RolloutSpec",
    "Side",
    "RUNTIME_SAFETY_MARGIN",
    "estimate_runtime_s",
@@ -26,7 +26,7 @@ from giant.analysis.reduce import (
    entry_axis,
    transverse_expr,
 )
-from giant.analysis.sources import Side, open_side, physical_steps, secondaries
+from giant.analysis.sources import RolloutSpec, Side, open_side, physical_steps, secondaries
 from giant.analysis.variables import RANGED_VARS


@@ -74,11 +74,9 @@ def _row_subsample(lf: pl.LazyFrame, sample_rows: int, seed: int) -> pl.LazyFram
    return lf.filter((pl.col("pre_E").hash(seed=seed) % 2**32) < threshold)


-def _combined_quantiles(
-    r_vals: np.ndarray, t_vals: np.ndarray, lo_q: float, hi_q: float
-) -> tuple[float, float]:
-    """Robust (lo_q, hi_q) range over the union of two value samples."""
-    both = np.concatenate([r_vals, t_vals])
+def _combined_quantiles(vals: list[np.ndarray], lo_q: float, hi_q: float) -> tuple[float, float]:
+    """Robust (lo_q, hi_q) range over the union of several value samples."""
+    both = np.concatenate(vals)
    lo, hi = float(np.quantile(both, lo_q)), float(np.quantile(both, hi_q))
    if not (hi - lo > 1e-6 * max(abs(hi), 1.0)):
        lo, hi = lo - 0.5, hi + 0.5
@@ -86,7 +84,7 @@ def _combined_quantiles(


 def build_context(
-    rollout: str | Path | pl.LazyFrame,
+    rollouts: list[RolloutSpec],
    reference: str | Path | pl.LazyFrame,
    *,
    n_energy_bins: int = 4,
@@ -96,58 +94,51 @@ def build_context(
    sample_rows: int = 1_000_000,
    seed: int = 0,
 ) -> Context:
-    """Resolve the shared context from the two files (the ``prep`` step)."""
-    r_all = open_side(rollout, Side.rollout)
+    """Resolve the shared context from the reference + every rollout (the ``prep`` step).
+
+    Every range/quantile below is the union of the reference and *all*
+    rollouts, so a single set of fixed bin edges/group sets is valid for
+    every series a compute job streams over.
+    """
    t_all = open_side(reference, Side.reference)
-    r_lf = physical_steps(r_all, Side.rollout)
    t_lf = physical_steps(t_all, Side.reference)
+    r_lfs = {rs.name: physical_steps(open_side(rs.source, Side.rollout), Side.rollout) for rs in rollouts}

    # Ranged marginal variables: robust ranges over a shared row subsample.
    exprs = [e.alias(n) for n, (_, e) in RANGED_VARS.items()]
-    r_s = (
-        _row_subsample(r_lf, sample_rows, seed)
-        .select(exprs)
-        .collect(engine="streaming")
-    )
-    t_s = (
-        _row_subsample(t_lf, sample_rows, seed)
-        .select(exprs)
-        .collect(engine="streaming")
-    )
+    t_s = _row_subsample(t_lf, sample_rows, seed).select(exprs).collect(engine="streaming")
+    r_s = {
+        name: _row_subsample(lf, sample_rows, seed).select(exprs).collect(engine="streaming")
+        for name, lf in r_lfs.items()
+    }
    var_ranges = {
-        name: _combined_quantiles(
-            r_s[name].to_numpy(), t_s[name].to_numpy(), _LO_Q, _HI_Q
-        )
+        name: _combined_quantiles([t_s[name].to_numpy(), *(df[name].to_numpy() for df in r_s.values())], _LO_Q, _HI_Q)
        for name in RANGED_VARS
    }

    # Energy-bin edges from exact per-event incident energies (cheap group_by).
    def _incident(lf: pl.LazyFrame) -> np.ndarray:
-        return (
-            lf.group_by("event_id")
-            .agg(pl.col("pre_E").max())
-            .collect(engine="streaming")["pre_E"]
-            .to_numpy()
-        )
+        return lf.group_by("event_id").agg(pl.col("pre_E").max()).collect(engine="streaming")["pre_E"].to_numpy()

-    r_inc, t_inc = _incident(r_lf), _incident(t_lf)
-    energy_edges = energy_bin_edges(np.concatenate([r_inc, t_inc]), n_energy_bins)
+    t_inc = _incident(t_lf)
+    r_inc = {name: _incident(lf) for name, lf in r_lfs.items()}
+    energy_edges = energy_bin_edges(np.concatenate([t_inc, *r_inc.values()]), n_energy_bins)

    # Top PDG species and material list (cheap single-column group_bys).
    def _counts(lf: pl.LazyFrame, col: str) -> pl.DataFrame:
        return lf.group_by(col).agg(pl.len().alias("n")).collect(engine="streaming")

    pdg_counts = (
-        pl.concat([_counts(r_lf, "pdg"), _counts(t_lf, "pdg")])
+        pl.concat([_counts(t_lf, "pdg"), *(_counts(lf, "pdg") for lf in r_lfs.values())])
        .group_by("pdg")
        .agg(pl.col("n").sum())
        .sort("n", descending=True)
    )
    top_pdgs = [int(x) for x in pdg_counts["pdg"].to_list()[:top_k_pdg]]
-    materials = sorted(
-        set(_counts(r_lf, "material")["material"].to_list())
-        | set(_counts(t_lf, "material")["material"].to_list())
-    )
+    material_set: set[str] = set(_counts(t_lf, "material")["material"].to_list())
+    for lf in r_lfs.values():
+        material_set |= set(_counts(lf, "material")["material"].to_list())
+    materials = sorted(material_set)

    # Shower depth / transverse ranges from a subsampled proxy.
    def _proxy(lf: pl.LazyFrame) -> tuple[np.ndarray, np.ndarray]:
@@ -159,19 +150,20 @@ def build_context(
        )
        return sub["d"].to_numpy(), sub["t"].to_numpy()

-    r_d, r_t = _proxy(r_lf)
    t_d, t_t = _proxy(t_lf)
-    d_lo, d_hi = _combined_quantiles(r_d, t_d, _LO_Q, _HI_Q)
+    r_proxy = {name: _proxy(lf) for name, lf in r_lfs.items()}
+    d_lo, d_hi = _combined_quantiles([t_d, *(p[0] for p in r_proxy.values())], _LO_Q, _HI_Q)
    depth_edges = np.linspace(d_lo, d_hi, n_marginal_bins + 1)
-    t_hi = max(float(np.quantile(np.concatenate([r_t, t_t]), _HI_Q)), 1e-6)
+    t_hi = max(float(np.quantile(np.concatenate([t_t, *(p[1] for p in r_proxy.values())]), _HI_Q)), 1e-6)
    transverse_edges = np.linspace(0.0, t_hi, n_marginal_bins + 1)

    # Secondary energy range.
-    r_se = secondaries(r_lf, Side.rollout).select("energy")
-    t_se = secondaries(t_all, Side.reference).select("energy")
-    r_se = _row_sample_col(r_se, sample_rows, seed)
-    t_se = _row_sample_col(t_se, sample_rows, seed)
-    sec_energy_range = _combined_quantiles(r_se, t_se, _LO_Q, _HI_Q)
+    t_se = _row_sample_col(secondaries(t_all, Side.reference).select("energy"), sample_rows, seed)
+    r_se = {
+        name: _row_sample_col(secondaries(lf, Side.rollout).select("energy"), sample_rows, seed)
+        for name, lf in r_lfs.items()
+    }
+    sec_energy_range = _combined_quantiles([t_se, *r_se.values()], _LO_Q, _HI_Q)

    return Context(
        n_marginal_bins=n_marginal_bins,
@@ -184,8 +176,8 @@ def build_context(
        sec_energy_range=sec_energy_range,
        n_sec_bins=n_sec_bins,
        n_events={
-            "rollout": len(r_inc),
            "reference": len(t_inc),
+            **{name: len(arr) for name, arr in r_inc.items()},
        },
    )

@@ -68,9 +68,7 @@ def energy_bin_edges(incident_E: np.ndarray, n_bins: int = 4) -> np.ndarray:

 def energy_bin_labels(edges: np.ndarray) -> list[str]:
    """``E in [lo, hi)`` labels for each bin defined by ``edges`` (MeV)."""
-    return [
-        f"E in [{edges[i]:.3g}, {edges[i + 1]:.3g}) MeV" for i in range(len(edges) - 1)
-    ]
+    return [f"E in [{edges[i]:.3g}, {edges[i + 1]:.3g}) MeV" for i in range(len(edges) - 1)]


 def digitize_expr(value: pl.Expr, edges: np.ndarray) -> pl.Expr:
@@ -88,9 +86,7 @@ def digitize_expr(value: pl.Expr, edges: np.ndarray) -> pl.Expr:
    return idx.clip(0, n_bins - 1)


-def event_energy_bins(
-    lf: pl.LazyFrame, edges: np.ndarray
-) -> tuple[np.ndarray, np.ndarray]:
+def event_energy_bins(lf: pl.LazyFrame, edges: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Per-event incident-energy bin: ``(event_ids, bin_idx)`` numpy arrays.

    Incident energy is ``max(pre_E)`` per event (the primary). One bounded
@@ -29,8 +29,17 @@ from giant.constants import TERM_ESCAPED


 def _bin_expr(value: pl.Expr, lo: float, hi: float, nbins: int) -> pl.Expr:
-    """Uniform bin index of ``value`` over ``[lo, hi]`` into ``nbins`` bins."""
-    return ((value - lo) / (hi - lo) * nbins).floor().cast(pl.Int32).clip(0, nbins - 1)
+    """Uniform bin index of ``value`` over ``[lo, hi]`` into ``nbins`` bins.
+
+    Out-of-range values clamp into the edge bins, and the clamp deliberately
+    happens in f64 *before* the integer cast: a rollout is free to emit a wildly
+    out-of-range outlier (a step_length of 1e10 mm, say) or an inf, whose
+    unclamped bin index overflows i32 and makes the cast fail outright. NaN has
+    no edge to clamp to, so it becomes null and is dropped by the callers below
+    — the same thing ``np.histogram`` does with it.
+    """
+    idx = ((value - lo) / (hi - lo) * nbins).floor().clip(0, nbins - 1)
+    return pl.when(idx.is_nan()).then(None).otherwise(idx).cast(pl.Int32)


 def hist1d(
@@ -50,6 +59,7 @@ def hist1d(
    group = pl.lit(0, dtype=pl.Int64) if group is None else group
    res = (
        lf.select(group.alias("_g"), _bin_expr(value, lo, hi, nbins).alias("_b"))
+        .drop_nulls("_b")
        .group_by("_g", "_b")
        .agg(pl.len().alias("_n"))
        .collect(engine="streaming")
@@ -142,9 +152,7 @@ def attach_entry_axis(lf: pl.LazyFrame, entry: pl.DataFrame) -> pl.LazyFrame:
    """
    ids = entry["event_id"].to_numpy()
    return lf.with_columns(
-        pl.col("event_id")
-        .replace_strict(ids, entry[col].to_numpy(), return_dtype=pl.Float64)
-        .alias(col)
+        pl.col("event_id").replace_strict(ids, entry[col].to_numpy(), return_dtype=pl.Float64).alias(col)
        for col in _ENTRY_AXIS_COLS
    )

@@ -190,6 +198,7 @@ def profile_partial(
            _bin_expr(coord, lo, hi, nbins).alias("_b"),
            weight.alias("_w"),
        )
+        .drop_nulls("_b")
        .group_by("event_id", "_b")
        .agg(pl.col("_w").sum().alias("_ws"))
        .collect(engine="streaming")
@@ -254,10 +263,7 @@ def leakage_fraction(lf: pl.LazyFrame) -> np.ndarray:
        lf.group_by("event_id")
        .agg(
            pl.col("edep").sum().alias("deposited"),
-            pl.col("pre_E")
-            .filter(pl.col("termination_reason") == TERM_ESCAPED)
-            .sum()
-            .alias("escaped"),
+            pl.col("pre_E").filter(pl.col("termination_reason") == TERM_ESCAPED).sum().alias("escaped"),
        )
        .collect(engine="streaming")
    )
@@ -265,3 +271,20 @@ def leakage_fraction(lf: pl.LazyFrame) -> np.ndarray:
    escaped = per_event["escaped"].fill_null(0.0).to_numpy()
    total = deposited + escaped
    return np.where(total > 0, escaped / total, 0.0)
+
+
+def sec_count_by_event(lf_all: pl.LazyFrame, sec_lf: pl.LazyFrame) -> tuple[np.ndarray, np.ndarray]:
+    """Per-event secondary count, zero-filled for events that produced none.
+
+    Two bounded per-event ``group_by``s — the full event set (from ``lf_all``)
+    and the secondary counts (from ``sec_lf``, see ``sources.secondaries``) —
+    merged in Python via a dict. Both results are event-granularity (not
+    per-row), so this stays in the same bounded-memory budget as
+    ``event_scalars``; a plain ``group_by`` on ``sec_lf`` alone would silently
+    drop zero-secondary events instead of zero-filling them.
+    """
+    ev = lf_all.select("event_id").unique().collect(engine="streaming")["event_id"].to_numpy()
+    cnt_df = sec_lf.group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")
+    cnt = dict(zip(cnt_df["event_id"].to_list(), cnt_df["n"].to_list()))
+    counts = np.array([cnt.get(int(e), 0) for e in ev], dtype=np.int64)
+    return ev, counts
@@ -11,15 +11,24 @@ import json
 from dataclasses import asdict, dataclass, field
 from pathlib import Path

-# Reduced.kind values:
-#   "overlay_hist"   rollout vs reference density histogram over shared edges
+# Reduced.kind values (payload keys a rollout series by name under
+# payload["series"], with the reference — where one exists — kept as one
+# distinguished payload["reference"] entry; see catalog.py's module
+# docstring for the full per-kind payload shape):
+#   "overlay_hist"   N-rollout-series vs reference density histogram over shared edges
 #   "grouped_hist"   one panel per group (energy/pdg/material), each an overlay
-#   "profile"        edep-weighted mean +/- event-RMS vs depth/radius, two series
-#   "bar"            per-category rollout vs reference bars (share / counts)
-#   "single_hist"    one series only (e.g. rollout leakage; reference has none)
-#   "router_gating"  stacked mean MoE gate weight vs energy, rollout + reference
-#   "router_share"   stacked bar of MoE top-1 dispatch share by category
-#   "unavailable"    plot not applicable to this run (e.g. non-MoE checkpoint)
+#   "profile"        edep-weighted mean +/- event-RMS vs depth/radius, N series + reference
+#   "bar"            per-category N-rollout-series vs reference bars (share / counts)
+#   "single_hist"    rollout-only series (e.g. leakage; reference has none)
+#   "router_gating"  stacked mean MoE gate weight vs energy, one rollout+reference
+#                    panel-pair per rollout with an enabled MoE router
+#   "router_share"   stacked bar of MoE top-1 dispatch share by category, one
+#                    panel per rollout with an enabled MoE router
+#   "router_specialization"  max gate weight vs energy (one scalar trend line
+#                    summarizing "router_gating"), per rollout with an enabled router
+#   "heatmap"        row x col matrix + colorbar, one panel per rollout (a
+#                    distance scorecard or a predicted-vs-true confusion matrix)
+#   "unavailable"    plot not applicable to this run (e.g. no MoE checkpoint)


@dataclass
@@ -9,10 +9,19 @@ streaming compute.
 For each reduced artifact it writes ``<out>/<family>/<id>.pdf`` plus a sibling
 ``<id>.yaml`` (per-plot gallery metadata) and a per-family ``metadata.yaml``.
 Optionally runs ``gallery generate`` to build the static HTML site.
+
+Every rollout series gets a stable color via ``ps.get_color(i)``, ``i`` being
+its position in ``payload["series"]`` — that position is fixed by the run's
+YAML/``--label`` order (threaded unchanged from ``condor.RunMeta.rollouts``
+through every ``PlotSpec``), so a given rollout keeps the same color across
+every plot in a run. The reference, where a plot has one, always draws in one
+fixed, distinct style (dark ink, dashed) instead of taking a slot in that
+cycle.
 """

 from __future__ import annotations

+import dataclasses
 import subprocess
 from pathlib import Path

@@ -22,7 +31,32 @@ import yaml

 from giant.analysis.reduced import Reduced

-_SERIES_LABELS = {"rollout": "rollout", "reference": "reference (Geant4)"}
+_REFERENCE_LABEL = "reference (Geant4)"
+
+_TEX_ESCAPE_MAP = {
+    "\\": r"\textbackslash{}",
+    "%": r"\%",
+    "&": r"\&",
+    "#": r"\#",
+    "$": r"\$",
+    "_": r"\_",
+    "{": r"\{",
+    "}": r"\}",
+}
+
+
+def _tex_escape(text: str) -> str:
+    """Escape characters LaTeX treats specially in catalog-authored title/xlabel
+    text (e.g. a literal ``%`` in a "90% of deposited energy" title, which
+    ``usetex`` otherwise reads as a comment marker and aborts the whole figure —
+    see gitea #81). A single pass over the *original* characters, so the
+    backslashes an escape itself introduces (e.g. ``\textbackslash{}``) are
+    never re-escaped."""
+    return "".join(_TEX_ESCAPE_MAP.get(c, c) for c in text)
+
+
+def _ref_color() -> str:
+    return ps.colors.INK["primary"]


 def _density(counts: list[int] | np.ndarray, edges: np.ndarray) -> np.ndarray:
@@ -33,10 +67,13 @@ def _density(counts: list[int] | np.ndarray, edges: np.ndarray) -> np.ndarray:
    return counts / (total * (edges[1] - edges[0]))


-def _overlay(ax, edges: np.ndarray, series: dict[str, list], log_y: bool) -> None:
-    for key in ("reference", "rollout"):
-        if key in series:
-            ax.stairs(_density(series[key], edges), edges, label=_SERIES_LABELS[key])
+def _overlay(ax, edges: np.ndarray, payload: dict, log_y: bool) -> None:
+    if "reference" in payload:
+        ax.stairs(
+            _density(payload["reference"], edges), edges, label=_REFERENCE_LABEL, color=_ref_color(), linestyle="--"
+        )
+    for i, (name, counts) in enumerate(payload.get("series", {}).items()):
+        ax.stairs(_density(counts, edges), edges, label=name, color=ps.get_color(i))
    if log_y:
        ax.set_yscale("log")

@@ -47,12 +84,12 @@ def _router_summary(router_cfg: dict) -> str:
    return f"{router_cfg.get('type', '?')}×{router_cfg.get('n_experts', '?')}"


-def _figure_params_v2(mc: dict, run_meta: dict) -> dict:
-    """`_figure_params` for a new-shape (nested) `model_config` — has a
+def _figure_params_v2(mc: dict, meta: dict) -> dict:
+    """`_figure_params_single` for a new-shape (nested) `model_config` — has a
    `stage1_model` key. Reports stage 1's architecture (the headline
    generator); stage 2's generator is only added (`mode_s2`) when it
-    differs from stage 1's, since a mixed run (docs/v0.3.0-design.md's
-    `stage1=flow` + `stage2=wgan` case) is the interesting exception, not
+    differs from stage 1's, since a mixed run (the `stage1=flow` +
+    `stage2=wgan` case) is the interesting exception, not
    the common case."""
    s1 = mc["stage1_model"]
    s2 = mc.get("stage2_model") or {}
@@ -71,23 +108,24 @@ def _figure_params_v2(mc: dict, run_meta: dict) -> dict:
    if particle_type is not None:
        params["conditioning"] = particle_type
    params["router"] = _router_summary(s1.get("router") or {})
-    if run_meta.get("training_epoch") is not None:
-        params["epoch"] = run_meta["training_epoch"]
-    if run_meta.get("best_val_loss") is not None:
-        params["best_val_loss"] = round(run_meta["best_val_loss"], 4)
+    if meta.get("training_epoch") is not None:
+        params["epoch"] = meta["training_epoch"]
+    if meta.get("best_val_loss") is not None:
+        params["best_val_loss"] = round(meta["best_val_loss"], 4)
    if mode == "wgan":
        noise_dim = (s1.get("wgan") or {}).get("noise_dim")
        if noise_dim is not None:
            params["noise_dim"] = noise_dim
-    elif run_meta.get("steps") is not None:
-        params["steps"] = run_meta["steps"]
+    elif meta.get("steps") is not None:
+        params["steps"] = meta["steps"]
    return params


-def _figure_params(run_meta: dict) -> dict:
-    """Curated run identity for the figure subtitle (``new_figure(params=...)``).
+def _figure_params_single(meta: dict) -> dict:
+    """Curated run identity for the figure subtitle (``new_figure(params=...)``),
+    for exactly one rollout's ``plot_meta``.

-    ``run_meta``/each plot's own ``<id>.yaml`` (see ``_plot_metadata``) already
+    ``meta``/each plot's own ``<id>.yaml`` (see ``_plot_metadata``) already
    carry every threaded model/training/rollout/dataset parameter for
    after-the-fact lookup — this picks only the handful that matter for
    telling figures apart at a glance while flipping through a gallery, since
@@ -99,9 +137,9 @@ def _figure_params(run_meta: dict) -> dict:
    Handles both a v0.2 checkpoint's flat ``model_config`` and a v0.3.0
    nested one (has a ``stage1_model`` key — see ``_figure_params_v2``).
    """
-    mc = run_meta.get("model_config") or {}
+    mc = meta.get("model_config") or {}
    if "stage1_model" in mc:
-        return _figure_params_v2(mc, run_meta)
+        return _figure_params_v2(mc, meta)

    mode = mc.get("mode")
    params: dict = {}
@@ -114,18 +152,35 @@ def _figure_params(run_meta: dict) -> dict:
    if mc.get("conditioning") is not None:
        params["conditioning"] = mc["conditioning"]
    params["router"] = _router_summary(mc.get("router") or {})
-    if run_meta.get("training_epoch") is not None:
-        params["epoch"] = run_meta["training_epoch"]
-    if run_meta.get("best_val_loss") is not None:
-        params["best_val_loss"] = round(run_meta["best_val_loss"], 4)
+    if meta.get("training_epoch") is not None:
+        params["epoch"] = meta["training_epoch"]
+    if meta.get("best_val_loss") is not None:
+        params["best_val_loss"] = round(meta["best_val_loss"], 4)
    if mode == "wgan":
        if mc.get("noise_dim") is not None:
            params["noise_dim"] = mc["noise_dim"]
-    elif run_meta.get("steps") is not None:
-        params["steps"] = run_meta["steps"]
+    elif meta.get("steps") is not None:
+        params["steps"] = meta["steps"]
    return params


+def _figure_params(run_meta: dict) -> dict:
+    """Curated run identity for the figure subtitle.
+
+    A single-rollout run reuses that rollout's ``plot_meta`` (same curated
+    model/training/rollout subset as always — see ``_figure_params_single``);
+    a multi-rollout run instead names the series being compared, since no
+    single ``model_config`` applies to the figure as a whole (each plot's own
+    gallery YAML still carries every rollout's full ``plot_meta`` for
+    after-the-fact lookup, via ``_plot_metadata``).
+    """
+    rollouts = run_meta.get("rollouts") or {}
+    if len(rollouts) == 1:
+        ((_, meta),) = rollouts.items()
+        return _figure_params_single(meta)
+    return {"rollouts": ", ".join(rollouts)} if rollouts else {}
+
+
 def _render_overlay(r: Reduced, params: dict):
    edges = np.asarray(r.payload["edges"])
    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
@@ -139,9 +194,8 @@ def _render_overlay(r: Reduced, params: dict):
 def _render_single(r: Reduced, params: dict):
    edges = np.asarray(r.payload["edges"])
    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
-    ax.stairs(
-        _density(r.payload["rollout"], edges), edges, label=_SERIES_LABELS["rollout"]
-    )
+    for i, (name, counts) in enumerate(r.payload.get("series", {}).items()):
+        ax.stairs(_density(counts, edges), edges, label=name, color=ps.get_color(i))
    if r.payload.get("log_y"):
        ax.set_yscale("log")
    if r.payload.get("log_x"):
@@ -183,13 +237,17 @@ def _render_profile(r: Reduced, params: dict):
    edges = np.asarray(r.payload["edges"])
    centers = 0.5 * (edges[:-1] + edges[1:])
    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
-    for key in ("reference", "rollout"):
-        mean = np.asarray(r.payload[f"{key}_mean"])
-        std = np.asarray(r.payload[f"{key}_std"])
-        (line,) = ax.plot(centers, mean, label=_SERIES_LABELS[key])
-        ax.fill_between(
-            centers, mean - std, mean + std, alpha=0.2, color=line.get_color()
-        )
+    if "reference" in r.payload:
+        ref = r.payload["reference"]
+        mean, std = np.asarray(ref["mean"]), np.asarray(ref["std"])
+        color = _ref_color()
+        ax.plot(centers, mean, label=_REFERENCE_LABEL, color=color, linestyle="--")
+        ax.fill_between(centers, mean - std, mean + std, alpha=0.2, color=color)
+    for i, (name, side) in enumerate(r.payload.get("series", {}).items()):
+        mean, std = np.asarray(side["mean"]), np.asarray(side["std"])
+        color = ps.get_color(i)
+        ax.plot(centers, mean, label=name, color=color)
+        ax.fill_between(centers, mean - std, mean + std, alpha=0.2, color=color)
    ax.set_xlabel(r.xlabel)
    ax.set_ylabel(r.payload.get("ylabel", "mean deposited energy [MeV]"))
    ps.style_legend(ax, title="source")
@@ -199,12 +257,19 @@ def _render_profile(r: Reduced, params: dict):
 def _render_bar(r: Reduced, params: dict):
    labels = r.payload["labels"]
    x = np.arange(len(labels))
-    width = 0.4
+    series = r.payload.get("series", {})
+    has_ref = "reference" in r.payload
+    n_bars = len(series) + (1 if has_ref else 0)
+    width = 0.8 / max(n_bars, 1)
+    offsets = np.linspace(-0.4 + width / 2, 0.4 - width / 2, n_bars)
    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
-    ax.bar(
-        x - width / 2, r.payload["reference"], width, label=_SERIES_LABELS["reference"]
-    )
-    ax.bar(x + width / 2, r.payload["rollout"], width, label=_SERIES_LABELS["rollout"])
+    idx = 0
+    if has_ref:
+        ax.bar(x + offsets[idx], r.payload["reference"], width, label=_REFERENCE_LABEL, color=_ref_color())
+        idx += 1
+    for i, (name, vals) in enumerate(series.items()):
+        ax.bar(x + offsets[idx], vals, width, label=name, color=ps.get_color(i))
+        idx += 1
    ax.set_xticks(x)
    ax.set_xticklabels(labels, rotation=45, ha="right")
    ax.set_ylabel(r.payload.get("ylabel", "value"))
@@ -213,60 +278,137 @@ def _render_bar(r: Reduced, params: dict):


 def _render_router_gating(r: Reduced, params: dict):
-    n_experts = r.payload["n_experts"]
+    series = r.payload.get("series", {})
+    names = list(series)
    log_x = r.payload.get("log_x", False)
-    fig, axes = ps.new_figure(
-        "slide-16x9", title=r.title, params=params, nrows=1, ncols=2, squeeze=False
-    )
-    flat = axes.ravel()
-    for ax, key in zip(flat, ("rollout", "reference")):
-        side = r.payload.get(key, {})
-        centers = np.asarray(side.get("centers", []))
-        means = np.asarray(side.get("means", []))
-        if len(centers) and means.size:
-            cum = np.zeros(len(centers))
-            for i in range(n_experts):
-                ax.fill_between(
-                    centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}"
-                )
-                cum = cum + means[:, i]
-        if log_x:
-            ax.set_xscale("log")
-        ax.set_ylim(0, 1)
-        ax.set_title(_SERIES_LABELS[key], fontsize=8)
-        ax.set_xlabel(r.xlabel)
-    flat[0].set_ylabel("mean gate weight")
-    ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router")
+    fig, axes = ps.new_figure("slide-16x9", title=r.title, params=params, nrows=len(names), ncols=2, squeeze=False)
+    for row, name in enumerate(names):
+        entry = series[name]
+        n_experts = entry["n_experts"]
+        for col, key in enumerate(("rollout", "reference")):
+            ax = axes[row, col]
+            side = entry.get(key, {})
+            centers = np.asarray(side.get("centers", []))
+            means = np.asarray(side.get("means", []))
+            if len(centers) and means.size:
+                cum = np.zeros(len(centers))
+                for i in range(n_experts):
+                    ax.fill_between(centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}")
+                    cum = cum + means[:, i]
+            if log_x:
+                ax.set_xscale("log")
+            ax.set_ylim(0, 1)
+            panel_label = _REFERENCE_LABEL if key == "reference" else "rollout"
+            ax.set_title(f"{name} — {panel_label}", fontsize=8)
+            if row == len(names) - 1:
+                ax.set_xlabel(r.xlabel)
+        axes[row, 0].set_ylabel("mean gate weight")
+    if names:
+        ps.style_legend(axes[0, 0], title=f"{series[names[0]]['router_type']} router")
    return fig


 def _render_router_share(r: Reduced, params: dict):
-    categories = r.payload["categories"]
-    n_experts = r.payload["n_experts"]
-    x = np.arange(len(categories))
-    present = [k for k in ("rollout", "reference") if k in r.payload]
+    series = r.payload.get("series", {})
+    names = list(series)
+    present: tuple[str, ...] = ("rollout", "reference")
+    if names:
+        present = tuple(k for k in ("rollout", "reference") if k in series[names[0]])
+    ncols = max(len(present), 1)
+    fig, axes = ps.new_figure("slide-16x9", title=r.title, params=params, nrows=len(names), ncols=ncols, squeeze=False)
+    for row, name in enumerate(names):
+        entry = series[name]
+        n_experts = entry["n_experts"]
+        cats = entry["categories"]
+        x = np.arange(len(cats))
+        for col, key in enumerate(present):
+            ax = axes[row, col]
+            side = entry.get(key)
+            if side is not None:
+                shares = np.array([side[c] for c in cats])  # (n_cat, n_experts)
+                bottom = np.zeros(len(cats))
+                for i in range(n_experts):
+                    ax.bar(x, shares[:, i], bottom=bottom, label=f"expert {i}")
+                    bottom += shares[:, i]
+            ax.set_xticks(x)
+            ax.set_xticklabels(cats, rotation=45, ha="right")
+            ax.set_ylim(0, 1)
+            panel_label = _REFERENCE_LABEL if key == "reference" else "rollout"
+            ax.set_title(f"{name} — {panel_label}", fontsize=8)
+        axes[row, 0].set_ylabel("share of rows dispatched to expert")
+    if names:
+        ps.style_legend(axes[0, 0], title=f"{series[names[0]]['router_type']} router")
+    return fig
+
+
+def _render_router_specialization(r: Reduced, params: dict):
+    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
+    series = r.payload.get("series", {})
+    chance_levels: set[float] = set()
+    for i, (name, entry) in enumerate(series.items()):
+        color = ps.get_color(i)
+        if entry.get("chance_level") is not None:
+            chance_levels.add(entry["chance_level"])
+        for key, linestyle, label in (
+            ("rollout", "-", name),
+            ("reference", "--", f"{name} ({_REFERENCE_LABEL})"),
+        ):
+            side = entry.get(key)
+            if side and side["centers"]:
+                ax.plot(
+                    side["centers"],
+                    side["score"],
+                    label=label,
+                    color=color,
+                    linestyle=linestyle,
+                    marker="o",
+                    markersize=3,
+                )
+    for lvl in sorted(chance_levels):
+        ax.axhline(lvl, linestyle=":", color="gray")
+    if r.payload.get("log_x"):
+        ax.set_xscale("log")
+    ax.set_ylim(0, 1)
+    ax.set_xlabel(r.xlabel)
+    ax.set_ylabel("max gate weight")
+    ps.style_legend(ax, title="router")
+    return fig
+
+
+def _render_heatmap(r: Reduced, params: dict):
+    series = r.payload["series"]
+    row_labels = r.payload["row_labels"]
+    col_labels = r.payload["col_labels"]
+    names = list(series)
    fig, axes = ps.new_figure(
-        "slide-16x9",
+        "slide-16x9" if len(names) > 1 else "thesis-single",
        title=r.title,
        params=params,
        nrows=1,
-        ncols=len(present),
+        ncols=len(names),
        squeeze=False,
    )
    flat = axes.ravel()
-    for ax, key in zip(flat, present):
-        side = r.payload[key]
-        shares = np.array([side[c] for c in categories])  # (n_cat, n_experts)
-        bottom = np.zeros(len(categories))
-        for i in range(n_experts):
-            ax.bar(x, shares[:, i], bottom=bottom, label=f"expert {i}")
-            bottom += shares[:, i]
-        ax.set_xticks(x)
-        ax.set_xticklabels(categories, rotation=45, ha="right")
-        ax.set_ylim(0, 1)
-        ax.set_title(_SERIES_LABELS[key], fontsize=8)
-    flat[0].set_ylabel("share of rows dispatched to expert")
-    ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router")
+    im = None
+    for ax, name in zip(flat, names):
+        mat = np.asarray(series[name], dtype=float)
+        im = ax.imshow(
+            mat,
+            origin="upper",
+            aspect="auto",
+            cmap=r.payload.get("cmap", "viridis"),
+            vmin=r.payload.get("vmin"),
+            vmax=r.payload.get("vmax"),
+        )
+        ax.set_xticks(range(len(col_labels)))
+        ax.set_xticklabels(col_labels, rotation=45, ha="right")
+        ax.set_yticks(range(len(row_labels)))
+        ax.set_yticklabels(row_labels)
+        ax.set_xlabel(r.xlabel)
+        if len(names) > 1:
+            ax.set_title(name, fontsize=8)
+    flat[0].set_ylabel(r.payload.get("ylabel", ""))
+    fig.colorbar(im, ax=list(flat), label=r.payload.get("cbar_label", "value"))
    return fig


@@ -294,13 +436,21 @@ _RENDERERS = {
    "bar": _render_bar,
    "router_gating": _render_router_gating,
    "router_share": _render_router_share,
+    "router_specialization": _render_router_specialization,
+    "heatmap": _render_heatmap,
    "unavailable": _render_unavailable,
 }


 def render(r: Reduced, run_meta: dict | None = None):
-    """Build the matplotlib figure for one reduced artifact (dispatch on kind)."""
-    return _RENDERERS[r.kind](r, _figure_params(run_meta or {}))
+    """Build the matplotlib figure for one reduced artifact (dispatch on kind).
+
+    ``title``/``xlabel`` are LaTeX-escaped here, at the one point every kind's
+    renderer draws them from — ``_plot_metadata`` deliberately keeps using the
+    unescaped ``r`` for the gallery YAML, which isn't LaTeX.
+    """
+    escaped = dataclasses.replace(r, title=_tex_escape(r.title), xlabel=_tex_escape(r.xlabel))
+    return _RENDERERS[r.kind](escaped, _figure_params(run_meta or {}))


 def _plot_metadata(r: Reduced, run_meta: dict) -> dict:
@@ -347,9 +497,7 @@ def render_all(
        families.add(r.family)
        fig = render(r, run_meta)
        ps.savefig(fig, str(family_dir / r.id), formats=("pdf",))
-        (family_dir / f"{r.id}.yaml").write_text(
-            yaml.safe_dump(_plot_metadata(r, run_meta), sort_keys=False)
-        )
+        (family_dir / f"{r.id}.yaml").write_text(yaml.safe_dump(_plot_metadata(r, run_meta), sort_keys=False))
        pdfs.append(family_dir / f"{r.id}.pdf")
        import matplotlib.pyplot as plt

@@ -361,7 +509,7 @@ def render_all(
        yaml.safe_dump(
            {
                "title": run_meta.get("title", "GIANT rollout analysis"),
-                "description": "Autoregressive rollout compared against held-out Geant4 reference steps.",
+                "description": "Autoregressive rollout(s) compared against a held-out Geant4 reference steps file.",
                "experiment": "GIANT",
                "parameters": {k: v for k, v in run_meta.items() if k != "title"},
            },
@@ -370,9 +518,7 @@ def render_all(
    )
    for fam in families:
        (out_dir / fam / "metadata.yaml").write_text(
-            yaml.safe_dump(
-                {"title": fam, "description": f"{fam} plots."}, sort_keys=False
-            )
+            yaml.safe_dump({"title": fam, "description": f"{fam} plots."}, sort_keys=False)
        )

    if run_gallery:
@@ -389,17 +535,14 @@ def render_run(run_dir: str | Path, *, run_gallery: bool = False) -> list[Path]:
    (checkpoint, paths, cutoffs) from ``run_meta.json`` into every plot's
    gallery metadata and renders.
    """
-    from giant.analysis.condor import RunMeta, merge_all
+    from giant.analysis.run import RunMeta, merge_all

    run_dir = Path(run_dir)
    merge_all(run_dir)
    meta = RunMeta.load(run_dir / "run_meta.json")
    run_meta = {
        "title": meta.title,
-        "rollout": meta.rollout,
        "reference": meta.reference,
-        **meta.plot_meta,
+        "rollouts": {ro["name"]: ro["plot_meta"] for ro in meta.rollouts},
    }
-    return render_all(
-        run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery
-    )
+    return render_all(run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery)
@@ -35,6 +35,7 @@ from giant.analysis.reduced import Reduced
 if TYPE_CHECKING:
    import torch

+    from giant.analysis.sources import RolloutSide
    from giant.data.transforms import Normalizer

 _SAMPLE_ROWS = 200_000
@@ -66,27 +67,11 @@ class _RouterHandle:
    router_type: str


-def _conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]:
-    """(particle_conditioning, material_conditioning) for
-    `giant.data.transforms.build_cond_features` — from either a v0.2
-    checkpoint's flat `model_config["conditioning"]` (one shared string, same
-    for both axes) or a new-format one (independent
-    `model_config["conditioning"]["particle"/"material"]["type"]` —
-    docs/v0.3.0-design.md §3.1: the two axes may differ). Mirrors
-    `giant.cli._conditioning_axes`."""
-    raw = model_cfg.get("conditioning", default)
-    if isinstance(raw, dict):
-        return (
-            raw.get("particle", {}).get("type", default),
-            raw.get("material", {}).get("type", default),
-        )
-    return raw, raw
-
-
 def load_router(checkpoint: str | Path) -> _RouterHandle | None:
    """Load a checkpoint's Stage-1 router, or None if it isn't a MoE checkpoint."""
    import torch

+    from giant.checkpoint_io import conditioning_axes
    from giant.data.transforms import Normalizer
    from giant.model.network import build_models

@@ -95,9 +80,7 @@ def load_router(checkpoint: str | Path) -> _RouterHandle | None:
    # New nested shape (has a "stage1_model" key) vs. a v0.2 checkpoint's
    # flat model_config.
    router_cfg = (
-        (model_cfg.get("stage1_model") or {}).get("router")
-        if "stage1_model" in model_cfg
-        else model_cfg.get("router")
+        (model_cfg.get("stage1_model") or {}).get("router") if "stage1_model" in model_cfg else model_cfg.get("router")
    )
    if not router_cfg or not router_cfg.get("enabled"):
        return None
@@ -112,7 +95,7 @@ def load_router(checkpoint: str | Path) -> _RouterHandle | None:
    if router is None:
        return None

-    particle_conditioning, material_conditioning = _conditioning_axes(model_cfg)
+    particle_conditioning, material_conditioning = conditioning_axes(model_cfg)
    return _RouterHandle(
        router=router,
        pdg_map={int(k): v for k, v in ckpt["pdg_map"].items()},
@@ -124,9 +107,7 @@ def load_router(checkpoint: str | Path) -> _RouterHandle | None:
    )


-def _subsample(
-    lf: pl.LazyFrame, n: int, seed: int, extra_cols: tuple = ()
-) -> pl.DataFrame:
+def _subsample(lf: pl.LazyFrame, n: int, seed: int, extra_cols: tuple = ()) -> pl.DataFrame:
    total = lf.select(pl.len()).collect(engine="streaming").item()
    if total > n:
        threshold = int(n / total * 2**32)
@@ -134,9 +115,7 @@ def _subsample(
    return lf.select(*_COLS, *extra_cols).collect(engine="streaming")


-def _gate_for_df(
-    handle: _RouterHandle, df: pl.DataFrame
-) -> tuple[pl.DataFrame, np.ndarray]:
+def _gate_for_df(handle: _RouterHandle, df: pl.DataFrame) -> tuple[pl.DataFrame, np.ndarray]:
    """(filtered df, gate_weights) for rows in ``df`` with a known pdg/material.

    Rows whose species or material never appeared in the checkpoint's
@@ -160,13 +139,9 @@ def _gate_for_df(
    df = df.filter(pl.Series(known, dtype=pl.Boolean))

    data = {
-        "pre_pos": np.column_stack(
-            [df["pre_x"].to_numpy(), df["pre_y"].to_numpy(), df["pre_z"].to_numpy()]
-        ),
+        "pre_pos": np.column_stack([df["pre_x"].to_numpy(), df["pre_y"].to_numpy(), df["pre_z"].to_numpy()]),
        "pre_E": df["pre_E"].to_numpy(),
-        "pre_dir": np.column_stack(
-            [df["pre_dx"].to_numpy(), df["pre_dy"].to_numpy(), df["pre_dz"].to_numpy()]
-        ),
+        "pre_dir": np.column_stack([df["pre_dx"].to_numpy(), df["pre_dy"].to_numpy(), df["pre_dz"].to_numpy()]),
        "layer_id": df["layer_id"].to_numpy(),
        "pdg": df["pdg"].to_numpy(),
        "material": df["material"].to_numpy(),
@@ -180,9 +155,7 @@ def _gate_for_df(
        material_conditioning=handle.material_conditioning,
    )
    with torch.no_grad():
-        gate = handle.router.gate(
-            torch.from_numpy(cond_cont).float(), torch.from_numpy(cond_cat).long()
-        ).numpy()
+        gate = handle.router.gate(torch.from_numpy(cond_cont).float(), torch.from_numpy(cond_cat).long()).numpy()
    return df, gate


@@ -205,9 +178,7 @@ def _quantile_bins(x: np.ndarray, gate: np.ndarray, n_bins: int) -> dict:
    return {"centers": centers[valid].tolist(), "means": means[valid].tolist()}


-def _top1_shares(
-    categories: np.ndarray, idx: np.ndarray, order: list, n_experts: int
-) -> dict[str, list[float]]:
+def _top1_shares(categories: np.ndarray, idx: np.ndarray, order: list, n_experts: int) -> dict[str, list[float]]:
    """Fraction of each category's rows hard-dispatched to each expert.

    Uses `Router.top1` (argmax), not the soft `gate` mean — grouped top-1
@@ -227,15 +198,13 @@ def _top1_shares(
    return shares


-_NOTE_NOT_MOE = (
-    "checkpoint has no enabled MoE router (model.router.enabled is "
-    "false/absent) — nothing to show"
-)
+_NOTE_NOT_MOE = "checkpoint has no enabled MoE router (model.router.enabled is false/absent) — nothing to show"

 _TITLES = {
    "router_gating": "Router gating (mixture-of-experts decision boundaries)",
    "router_share_by_pdg": "Router expert share by particle species",
    "router_share_by_process": "Router expert share by physics process",
+    "router_specialization": "Router specialization score vs energy (max gate weight)",
 }


@@ -250,53 +219,103 @@ def _unavailable(spec_id: str) -> Reduced:
    )


-def compute_router_gating(
-    checkpoint: str | Path | None,
-    r_phys: pl.LazyFrame,
-    t_phys: pl.LazyFrame,
-    seed: int = 0,
-) -> Reduced:
-    """`Reduced` for the router-gating figure, or an explanatory note if n/a."""
+def _gating_entry(checkpoint: str | Path | None, r_phys: pl.LazyFrame, t_phys: pl.LazyFrame, seed: int) -> dict | None:
+    """One rollout's ``router_gating`` panel data, or ``None`` if not a MoE checkpoint."""
    handle = load_router(checkpoint) if checkpoint else None
    if handle is None:
-        return _unavailable("router_gating")
-
+        return None
    sides: dict[str, dict] = {}
    for name, lf in (("rollout", r_phys), ("reference", t_phys)):
        df = _subsample(lf, _SAMPLE_ROWS, seed)
        df, gate = _gate_for_df(handle, df)
        x = df["pre_E"].to_numpy()
-        sides[name] = (
-            _quantile_bins(x, gate, _N_BINS) if len(x) else {"centers": [], "means": []}
-        )
+        sides[name] = _quantile_bins(x, gate, _N_BINS) if len(x) else {"centers": [], "means": []}
+    return {"router_type": handle.router_type, "n_experts": handle.router.n_experts, **sides}

+
+def compute_router_gating(rollouts: dict[str, "RolloutSide"], t_phys: pl.LazyFrame, seed: int = 0) -> Reduced:
+    """`Reduced` for the router-gating figure: one panel-pair per rollout with
+    an enabled MoE router, or an explanatory note if none of them have one."""
+    series = {}
+    for name, rs in rollouts.items():
+        entry = _gating_entry(rs.checkpoint, rs.phys, t_phys, seed)
+        if entry is not None:
+            series[name] = entry
+    if not series:
+        return _unavailable("router_gating")
    return Reduced(
        id="router_gating",
        family="model",
        kind="router_gating",
        title=_TITLES["router_gating"],
        xlabel="pre-step energy [MeV]",
-        payload={
-            "router_type": handle.router_type,
-            "n_experts": handle.router.n_experts,
-            "log_x": True,
-            **sides,
-        },
+        payload={"series": series, "log_x": True},
    )


-def compute_router_share_by_pdg(
-    checkpoint: str | Path | None,
-    r_phys: pl.LazyFrame,
-    t_phys: pl.LazyFrame,
-    top_pdgs: list[int],
-    seed: int = 0,
-) -> Reduced:
-    """Stacked-bar share of each particle species dispatched to each expert."""
+def _specialization_entry(
+    checkpoint: str | Path | None, r_phys: pl.LazyFrame, t_phys: pl.LazyFrame, seed: int
+) -> dict | None:
+    """One rollout's ``router_specialization`` curve data, or ``None`` if not a MoE checkpoint.
+
+    Scalar specialization trend: max gate weight vs energy, per side.
+    Summarizes `router_gating`'s full per-expert stacked area into one curve —
+    the routing plan's own "how sharp is the boundary here" number (1/n_experts
+    = uniform/no specialization, 1.0 = one expert fully owns that energy). Same
+    quantile energy bins as `router_gating` (`_quantile_bins`), so this is
+    directly comparable to that plot's ceiling described in the roadmap's MoE
+    writeup.
+    """
    handle = load_router(checkpoint) if checkpoint else None
    if handle is None:
-        return _unavailable("router_share_by_pdg")
+        return None
+    sides: dict[str, dict] = {}
+    for name, lf in (("rollout", r_phys), ("reference", t_phys)):
+        df = _subsample(lf, _SAMPLE_ROWS, seed)
+        df, gate = _gate_for_df(handle, df)
+        x = df["pre_E"].to_numpy()
+        if len(x):
+            binned = _quantile_bins(x, gate, _N_BINS)
+            means = np.asarray(binned["means"])
+            score = means.max(axis=1).tolist() if means.size else []
+            sides[name] = {"centers": binned["centers"], "score": score}
+        else:
+            sides[name] = {"centers": [], "score": []}
+    return {
+        "router_type": handle.router_type,
+        "n_experts": handle.router.n_experts,
+        "chance_level": 1.0 / handle.router.n_experts,
+        **sides,
+    }

+
+def compute_router_specialization(rollouts: dict[str, "RolloutSide"], t_phys: pl.LazyFrame, seed: int = 0) -> Reduced:
+    """`Reduced` for the router-specialization figure, one curve per rollout with
+    an enabled MoE router (see `_specialization_entry`)."""
+    series = {}
+    for name, rs in rollouts.items():
+        entry = _specialization_entry(rs.checkpoint, rs.phys, t_phys, seed)
+        if entry is not None:
+            series[name] = entry
+    if not series:
+        return _unavailable("router_specialization")
+    return Reduced(
+        id="router_specialization",
+        family="model",
+        kind="router_specialization",
+        title=_TITLES["router_specialization"],
+        xlabel="pre-step energy [MeV]",
+        payload={"series": series, "log_x": True},
+    )
+
+
+def _share_by_pdg_entry(
+    checkpoint: str | Path | None, r_phys: pl.LazyFrame, t_phys: pl.LazyFrame, top_pdgs: list[int], seed: int
+) -> dict | None:
+    """One rollout's ``router_share_by_pdg`` panel-pair data, or ``None`` if not a MoE checkpoint."""
+    handle = load_router(checkpoint) if checkpoint else None
+    if handle is None:
+        return None
    labels = [pdg_label(p) for p in top_pdgs]
    sides: dict[str, dict] = {}
    for name, lf in (("rollout", r_phys), ("reference", t_phys)):
@@ -304,69 +323,82 @@ def compute_router_share_by_pdg(
        df, gate = _gate_for_df(handle, df)
        if len(df):
            idx = gate.argmax(axis=1)
-            shares = _top1_shares(
-                df["pdg"].to_numpy(), idx, top_pdgs, handle.router.n_experts
-            )
+            shares = _top1_shares(df["pdg"].to_numpy(), idx, top_pdgs, handle.router.n_experts)
        else:
            shares = {str(p): [0.0] * handle.router.n_experts for p in top_pdgs}
        sides[name] = {labels[i]: shares[str(p)] for i, p in enumerate(top_pdgs)}
+    return {"router_type": handle.router_type, "n_experts": handle.router.n_experts, "categories": labels, **sides}

+
+def compute_router_share_by_pdg(
+    rollouts: dict[str, "RolloutSide"], t_phys: pl.LazyFrame, top_pdgs: list[int], seed: int = 0
+) -> Reduced:
+    """`Reduced` for the router expert-share-by-species figure, one panel-pair
+    per rollout with an enabled MoE router."""
+    series = {}
+    for name, rs in rollouts.items():
+        entry = _share_by_pdg_entry(rs.checkpoint, rs.phys, t_phys, top_pdgs, seed)
+        if entry is not None:
+            series[name] = entry
+    if not series:
+        return _unavailable("router_share_by_pdg")
    return Reduced(
        id="router_share_by_pdg",
        family="model",
        kind="router_share",
        title=_TITLES["router_share_by_pdg"],
        xlabel="particle species",
-        payload={
-            "router_type": handle.router_type,
-            "n_experts": handle.router.n_experts,
-            "categories": labels,
-            **sides,
-        },
+        payload={"series": series},
    )


-def compute_router_share_by_process(
-    checkpoint: str | Path | None,
-    t_phys: pl.LazyFrame,
-    seed: int = 0,
-    top_k: int = _TOP_K_PROCESS,
-) -> Reduced:
-    """Stacked-bar share of each physics process dispatched to each expert.
-
-    Reference-only: ``process`` is the true post-step physics process — a
-    label the rollout side has no equivalent of (see
-    `giant.model.network.ProcessRouter`, which predicts it from pre-step
-    conditioning alone, never observes it at eval time). This plot instead
-    checks *after the fact*, on real data, how well the router's conditioning
-    -based dispatch lines up with the true process.
-    """
+def _share_by_process_entry(checkpoint: str | Path | None, t_phys: pl.LazyFrame, seed: int, top_k: int) -> dict | None:
+    """One rollout checkpoint's ``router_share_by_process`` panel data (reference-only), or ``None`` if not MoE."""
    handle = load_router(checkpoint) if checkpoint else None
    if handle is None:
-        return _unavailable("router_share_by_process")
-
+        return None
    df = _subsample(t_phys, _SAMPLE_ROWS, seed, extra_cols=("process",))
    df, gate = _gate_for_df(handle, df)
    if len(df):
        counts = df["process"].value_counts().sort("count", descending=True)
        order = counts["process"].to_list()[:top_k]
        idx = gate.argmax(axis=1)
-        shares = _top1_shares(
-            df["process"].to_numpy(), idx, order, handle.router.n_experts
-        )
+        shares = _top1_shares(df["process"].to_numpy(), idx, order, handle.router.n_experts)
    else:
        order, shares = [], {}
+    return {
+        "router_type": handle.router_type,
+        "n_experts": handle.router.n_experts,
+        "categories": order,
+        "reference": {p: shares[p] for p in order},
+    }

+
+def compute_router_share_by_process(
+    rollouts: dict[str, "RolloutSide"], t_phys: pl.LazyFrame, seed: int = 0, top_k: int = _TOP_K_PROCESS
+) -> Reduced:
+    """Stacked-bar share of each physics process dispatched to each expert, one
+    panel per rollout checkpoint with an enabled MoE router.
+
+    Reference-only: ``process`` is the true post-step physics process — a
+    label the rollout side has no equivalent of (see
+    `giant.model.network.ProcessRouter`, which predicts it from pre-step
+    conditioning alone, never observes it at eval time). This plot instead
+    checks *after the fact*, on real data, how well each checkpoint's router
+    -based dispatch lines up with the true process.
+    """
+    series = {}
+    for name, rs in rollouts.items():
+        entry = _share_by_process_entry(rs.checkpoint, t_phys, seed, top_k)
+        if entry is not None:
+            series[name] = entry
+    if not series:
+        return _unavailable("router_share_by_process")
    return Reduced(
        id="router_share_by_process",
        family="model",
        kind="router_share",
        title=_TITLES["router_share_by_process"],
        xlabel="physics process",
-        payload={
-            "router_type": handle.router_type,
-            "n_experts": handle.router.n_experts,
-            "categories": order,
-            "reference": {p: shares[p] for p in order},
-        },
+        payload={"series": series},
    )
@@ -1,4 +1,6 @@
-"""HTCondor orchestration driven by a ``giant rollout`` YAML sidecar.
+"""Analysis run directories: prep, per-(plot, chunk) compute, and merge.
+
+Driven by one or more ``giant rollout`` YAML sidecars.

 A rollout writes a YAML sidecar (``giant/cli.py:_write_prediction_ref`` +
 rollout extras) that already names both files we need and carries the run's
@@ -10,8 +12,11 @@ provenance:
 * ``checkpoint``, ``geometry_oracle``, ``energy_cutoff``, ``steps``, ... —
  metadata that flows straight into every plot's gallery ``metadata.yaml``.

-So the analysis takes that one YAML as input, derives its own **run directory**
-next to the rollout parquet, and lays everything out under it:
+The analysis takes N such YAMLs — one series per rollout, all required to
+share the same ``dataset`` (the premise is "N candidates vs one ground
+truth") — resolves each one's series name (``load_rollout_yamls``), derives
+its own **run directory** next to the first rollout's parquet, and lays
+everything out under it:

    <run_dir>/shared.json     fixed bin edges / group sets (prep)
    <run_dir>/run_meta.json   resolved rollout/reference paths + plot metadata
@@ -19,13 +24,16 @@ next to the rollout parquet, and lays everything out under it:
    <run_dir>/reduced/<id>.json   merged, per plot
    <run_dir>/plots/<family>/<id>.pdf   rendered locally

-Job model (one condor job per (plot, chunk), compute/merge/render split):
+Job model (one job per (plot, chunk), compute/merge/render split). Job
+submission itself is b2luigi's (``giant/workflow/tasks.py`` — ``AnalysisPrepTask``
+/ ``AnalysisComputeTask`` / ``AnalysisRenderTask``); this module only provides
+the three steps they call:

-1. ``prep`` runs once on the submit node — reads the YAML, resolves the shared
+1. ``prep`` runs once locally — reads the YAML, resolves the shared
   context from a subsample, writes ``shared.json`` + ``run_meta.json``
   (including the run's configured ``n_chunks``).
 2. one job per catalog id x chunk index runs ``giant analyze compute-one
-   --run-dir`` on a worker — a single streaming pass over that
+   --run-dir`` (or ``compute_one`` in-process) on a worker — a single streaming pass over that
   ``event_id``-disjoint chunk, writing ``reduced_partial/<id>__<chunk>.json``
   (polars/numpy only, no LaTeX). Specs marked ``chunkable=False``
   (``PlotSpec``, ``catalog.py``) always run as a single chunk.
@@ -35,15 +43,16 @@ Job model (one condor job per (plot, chunk), compute/merge/render split):
   ``reduced/<id>.json``, then renders those into the styled PDF + gallery tree
   (that step imports plotstyle/LaTeX).

-Files on ``/ceph`` or ``/work`` are reached via ``ProvidesETPResources``; no
-HTCondor file transfer of the multi-GB inputs.
+Files on ``/ceph`` or ``/work`` are reached directly (see
+``giant/workflow/htcondor.py``); no HTCondor file transfer of the multi-GB
+inputs.
 """

 from __future__ import annotations

 import json
 import shutil
-import sys
+from collections.abc import Sequence
 from dataclasses import dataclass, field
 from pathlib import Path

@@ -53,8 +62,7 @@ import yaml
 from giant.analysis.catalog import Bundle, catalog_ids, get_spec
 from giant.analysis.context import Context, build_context
 from giant.analysis.reduced import Partial
-from giant.analysis.runtime_estimate import estimate_runtime_s
-from giant.analysis.sources import Side, open_side
+from giant.analysis.sources import RolloutSpec, Side, open_side

 # Keys copied verbatim from a rollout YAML into each plot's gallery metadata.
 _PLOT_META_KEYS = (
@@ -83,7 +91,7 @@ _PLOT_META_KEYS = (
    "best_val_loss",
    "training_config",
    "training_meta",
-    # §11.3 diagnostic — only present when giant rollout ran under
+    # Diagnostic — only present when giant rollout ran under
    # stage2_model.particle_type.target="embedding" (see giant/cli.py's
    # rollout command and giant.rollout.L1DistCollector); absent otherwise,
    # which the type_embedding_l1_distance PlotSpec (catalog.py) reads as
@@ -111,8 +119,60 @@ def load_rollout_yaml(path: str | Path) -> dict:
    return d


+@dataclass
+class LoadedRollout:
+    """One rollout YAML plus its resolved series ``name`` (see ``load_rollout_yamls``)."""
+
+    name: str
+    yaml: dict
+
+
+def load_rollout_yamls(
+    paths: Sequence[str | Path], labels: Sequence[str] | None = None
+) -> tuple[list[LoadedRollout], str]:
+    """Load every rollout YAML, resolve each one's series name, and verify they
+    all share one reference (``dataset``) file — the premise is "N candidates
+    vs one ground truth", not N independent comparisons.
+
+    Names: an explicit ``labels[i]`` if given (``labels`` must be empty or
+    exactly ``len(paths)`` long); otherwise the YAML's stem for N>1, or
+    ``"rollout"`` for the single-YAML case — matching today's one-series
+    legend/payload key, so a single-rollout run renders identically to
+    before this feature existed. Raises ``ValueError`` if two rollouts
+    resolve to the same name, or if the YAMLs don't all name the same
+    ``dataset``.
+    """
+    if labels and len(labels) != len(paths):
+        raise ValueError(f"--label given {len(labels)} time(s) but {len(paths)} rollout YAML(s) were passed")
+    yamls = [load_rollout_yaml(p) for p in paths]
+    if labels:
+        names = list(labels)
+    elif len(paths) == 1:
+        names = ["rollout"]
+    else:
+        names = [Path(p).stem for p in paths]
+    if len(set(names)) != len(names):
+        dupes = sorted({n for n in names if names.count(n) > 1})
+        raise ValueError(f"rollout series names collide: {dupes} — pass --label to disambiguate")
+
+    references = {str(y["dataset"]) for y in yamls}
+    if len(references) > 1:
+        detail = "\n".join(f"  {p}: dataset={y['dataset']!r}" for p, y in zip(paths, yamls))
+        raise ValueError(
+            "all rollout YAMLs must be seeded from the same reference (dataset) "
+            f"file — got {len(references)} distinct ones:\n{detail}"
+        )
+
+    return [LoadedRollout(name=n, yaml=y) for n, y in zip(names, yamls)], yamls[0]["dataset"]
+
+
+def _run_tag(y: dict) -> str:
+    rollout = Path(y["output"])
+    return str(y.get("prediction_id") or rollout.stem)[:8]
+
+
 def derive_run_dir(
-    rollout_yaml: dict,
+    rollout_yamls: list[dict],
    run_dir: str | Path | None = None,
    default_base: str | Path | None = None,
 ) -> Path:
@@ -122,14 +182,24 @@ def derive_run_dir(
    ``default_base / analysis_<tag>`` if ``default_base`` is given (the CLI
    passes the repo's gitignored ``analysis_runs/``, so run directories don't
    pile up on ``/ceph`` next to the rollout parquet). Falls back to next to
-    the rollout parquet — the original convention — for callers that don't
-    care where the run directory lives.
+    the *first* rollout's parquet — the original convention — for callers
+    that don't care where the run directory lives.
+
+    ``tag`` is a single rollout's ``prediction_id``/output stem (matching
+    today's single-rollout convention exactly) when there's only one; for
+    N>1 it joins up to three tags with ``-``, then ``-plus<K>`` for any
+    beyond that, so a many-rollout run still gets a short, stable directory
+    name.
    """
    if run_dir is not None:
        return Path(run_dir)
-    rollout = Path(rollout_yaml["output"])
-    tag = str(rollout_yaml.get("prediction_id") or rollout.stem)[:8]
-    base = Path(default_base) if default_base is not None else rollout.parent
+    tags = [_run_tag(y) for y in rollout_yamls]
+    if len(tags) == 1:
+        tag = tags[0]
+    else:
+        shown, rest = tags[:3], tags[3:]
+        tag = "-".join(shown) + (f"-plus{len(rest)}" if rest else "")
+    base = Path(default_base) if default_base is not None else Path(rollout_yamls[0]["output"]).parent
    return base / f"analysis_{tag}"


@@ -139,17 +209,21 @@ def _plot_meta(rollout_yaml: dict) -> dict:

@dataclass
 class RunMeta:
-    """Resolved paths + plot metadata for one analysis run (``run_meta.json``)."""
+    """Resolved paths + plot metadata for one analysis run (``run_meta.json``).

-    rollout: str
+    ``rollouts`` is ``[{"name", "path", "plot_meta"}, ...]``, insertion order
+    = the order rollouts were given on the CLI (and so the order every
+    ``Reduced.payload["series"]`` dict is built in — see ``catalog.py``).
+    """
+
+    rollouts: list[dict]
    reference: str
    run_dir: str
    title: str
-    plot_meta: dict
    n_chunks: int = 1
-    # rollout+reference row count of each event_id-disjoint chunk, and the
-    # dataset total — inputs to `runtime_estimate.estimate_runtime_s`. Empty/0
-    # on run directories written before this field existed.
+    # combined rollout+reference row count of each event_id-disjoint chunk,
+    # and the dataset total — inputs to `runtime_estimate.estimate_runtime_s`.
+    # Empty/0 on run directories written before this field existed.
    rows_per_chunk: list[int] = field(default_factory=list)
    total_rows: int = 0

@@ -161,10 +235,8 @@ class RunMeta:
        return cls(**json.loads(Path(path).read_text()))


-def _rows_per_chunk(
-    rollout: str | Path, reference: str | Path, n_chunks: int
-) -> list[int]:
-    """Rollout+reference row count of each ``event_id % n_chunks`` chunk.
+def _rows_per_chunk(rollouts: list[str | Path], reference: str | Path, n_chunks: int) -> list[int]:
+    """Combined rollout+reference row count of each ``event_id % n_chunks`` chunk.

    One cheap streaming ``group_by`` per side (just the ``event_id`` column) —
    the sizing input every job's estimated walltime
@@ -180,7 +252,8 @@ def _rows_per_chunk(
        )

    out = [0] * n_chunks
-    for lf in (open_side(rollout, Side.rollout), open_side(reference, Side.reference)):
+    sides = [open_side(reference, Side.reference)] + [open_side(r, Side.rollout) for r in rollouts]
+    for lf in sides:
        df = counts(lf)
        for c, n in zip(df["_c"].to_list(), df["n"].to_list()):
            out[c] += n
@@ -188,20 +261,22 @@ def _rows_per_chunk(


 def prep(
-    rollout_yaml: str | Path,
+    rollout_yamls: Sequence[str | Path],
    run_dir: str | Path | None = None,
    n_chunks: int = 1,
    default_base: str | Path | None = None,
+    labels: Sequence[str] | None = None,
    **ctx_kwargs,
 ) -> Path:
-    """Read the rollout YAML, build the shared context, and lay out the run dir.
+    """Read the rollout YAML(s), build the shared context, and lay out the run dir.

    Writes ``shared.json`` + ``run_meta.json`` and returns the run directory.
    ``n_chunks`` is the run-level chunk count every ``compute-one``/``merge-one``
    job reads back out of ``run_meta.json`` (via ``RunMeta.n_chunks``), so it is
    resolved once here rather than re-passed (and risking disagreement) at every
    later step. See ``derive_run_dir`` for how ``run_dir``/``default_base``
-    resolve the actual directory.
+    resolve the actual directory, and ``load_rollout_yamls`` for how
+    ``labels``/YAML stems resolve each rollout's series name.

    Clears any existing ``reduced_partial/``/``reduced/`` from a prior prep of
    this same ``run_dir``: partial files carry no record of what context
@@ -210,8 +285,8 @@ def prep(
    rollout/reference files changed) would otherwise let ``merge_one`` silently
    merge stale partials against the new ``shared.json``.
    """
-    y = load_rollout_yaml(rollout_yaml)
-    run_path = derive_run_dir(y, run_dir, default_base=default_base)
+    loaded, reference = load_rollout_yamls(list(rollout_yamls), labels)
+    run_path = derive_run_dir([lr.yaml for lr in loaded], run_dir, default_base=default_base)
    run_path.mkdir(parents=True, exist_ok=True)

    for stale in ("reduced_partial", "reduced"):
@@ -219,19 +294,22 @@ def prep(
        if stale_dir.exists():
            shutil.rmtree(stale_dir)

-    rollout, reference = y["output"], y["dataset"]
-    ctx = build_context(rollout, reference, **ctx_kwargs)
+    rollout_specs = [RolloutSpec(name=lr.name, source=lr.yaml["output"]) for lr in loaded]
+    ctx = build_context(rollout_specs, reference, **ctx_kwargs)
    ctx.save(run_path / "shared.json")

-    rows_per_chunk = _rows_per_chunk(rollout, reference, n_chunks)
+    rows_per_chunk = _rows_per_chunk([lr.yaml["output"] for lr in loaded], reference, n_chunks)
+
+    rollouts_meta = [
+        {"name": lr.name, "path": str(lr.yaml["output"]), "plot_meta": _plot_meta(lr.yaml)} for lr in loaded
+    ]
+    ckpts = ", ".join(Path(lr.yaml.get("checkpoint", "")).name or "rollout" for lr in loaded)

-    ckpt = Path(y.get("checkpoint", "")).name or "rollout"
    RunMeta(
-        rollout=str(rollout),
+        rollouts=rollouts_meta,
        reference=str(reference),
        run_dir=str(run_path),
-        title=f"GIANT rollout analysis — {ckpt}",
-        plot_meta=_plot_meta(y),
+        title=f"GIANT rollout analysis — {ckpts}",
        n_chunks=n_chunks,
        rows_per_chunk=rows_per_chunk,
        total_rows=sum(rows_per_chunk),
@@ -246,17 +324,19 @@ def prep(

 def compute_reduced(
    spec_id: str,
-    rollout: str | Path,
+    rollouts: list[dict],
    reference: str | Path,
    shared: str | Path,
    out: str | Path,
-    checkpoint: str | None = None,
    chunk_index: int = 0,
    n_chunks: int = 1,
-    type_embedding_l1_dist: dict | None = None,
 ) -> Path:
    """Core: run one (plot, chunk)'s partial reduction against explicit paths.

+    ``rollouts``: ``[{"name", "path", "checkpoint"?, "type_embedding_l1_dist"?},
+    ...]``, one per rollout series (insertion order preserved through to every
+    plot's ``Reduced.payload["series"]``).
+
    Writes a ``Partial`` JSON — the raw, not-yet-merged output of
    ``PlotSpec.compute_partial`` — never a finished ``Reduced``; ``merge_one``
    is what combines every chunk's ``Partial`` for a plot into the final
@@ -268,17 +348,18 @@ def compute_reduced(
    effective_n = n_chunks if spec.chunkable else 1
    if not (0 <= chunk_index < effective_n):
        raise ValueError(
-            f"{spec_id}: chunk_index={chunk_index} out of range for "
-            f"n_chunks={effective_n} (chunkable={spec.chunkable})"
+            f"{spec_id}: chunk_index={chunk_index} out of range for n_chunks={effective_n} (chunkable={spec.chunkable})"
        )
-    bundle = Bundle.open(
-        rollout,
-        reference,
-        ctx,
-        checkpoint=checkpoint,
-        chunk=(chunk_index, effective_n),
-        type_embedding_l1_dist=type_embedding_l1_dist,
-    )
+    rollout_specs = [
+        RolloutSpec(
+            name=r["name"],
+            source=r["path"],
+            checkpoint=r.get("checkpoint"),
+            type_embedding_l1_dist=r.get("type_embedding_l1_dist"),
+        )
+        for r in rollouts
+    ]
+    bundle = Bundle.open(rollout_specs, reference, ctx, chunk=(chunk_index, effective_n))
    partial = Partial(
        id=spec_id,
        family=spec.family,
@@ -294,16 +375,23 @@ def compute_one(spec_id: str, run_dir: str | Path, chunk_index: int = 0) -> Path
    """Run one (plot, chunk)'s partial reduction from a prepped run directory."""
    run_path = Path(run_dir)
    meta = RunMeta.load(run_path / "run_meta.json")
+    rollouts = [
+        {
+            "name": ro["name"],
+            "path": ro["path"],
+            "checkpoint": ro["plot_meta"].get("checkpoint"),
+            "type_embedding_l1_dist": ro["plot_meta"].get("type_embedding_l1_dist"),
+        }
+        for ro in meta.rollouts
+    ]
    return compute_reduced(
        spec_id,
-        meta.rollout,
+        rollouts,
        meta.reference,
        run_path / "shared.json",
        run_path / "reduced_partial" / f"{spec_id}__{chunk_index}.json",
-        checkpoint=meta.plot_meta.get("checkpoint"),
        chunk_index=chunk_index,
        n_chunks=meta.n_chunks,
-        type_embedding_l1_dist=meta.plot_meta.get("type_embedding_l1_dist"),
    )


@@ -327,10 +415,7 @@ def merge_one(spec_id: str, run_dir: str | Path) -> Path:
    effective_n = meta.n_chunks if spec.chunkable else 1

    partial_dir = run_path / "reduced_partial"
-    found = {
-        p.chunk: p
-        for p in (Partial.load(jf) for jf in partial_dir.glob(f"{spec_id}__*.json"))
-    }
+    found = {p.chunk: p for p in (Partial.load(jf) for jf in partial_dir.glob(f"{spec_id}__*.json"))}
    missing = sorted(set(range(effective_n)) - set(found))
    if missing:
        raise FileNotFoundError(
@@ -348,144 +433,3 @@ def merge_one(spec_id: str, run_dir: str | Path) -> Path:
 def merge_all(run_dir: str | Path) -> list[Path]:
    """Merge every catalog plot's chunk partials into ``reduced/<id>.json``."""
    return [merge_one(spec_id, run_dir) for spec_id in catalog_ids()]
-
-
-# ---------------------------------------------------------------------------
-# submit description
-# ---------------------------------------------------------------------------
-
-
-@dataclass
-class SubmitConfig:
-    run_dir: Path
-    accounting_group: str
-    repo_dir: Path
-    docker_image: str = "cverstege/alma9-gridjob"
-    request_memory_mb: int = 8192
-    request_cpus: int = 1
-    remote: bool = False  # +RemoteJob (grid I/O) vs ProvidesETPResources (local files)
-    n_chunks: int = 1  # per-plot data chunks; ignored for chunkable=False specs
-
-
-_WRAPPER = """#!/bin/bash
-set -euo pipefail
-cd {repo_dir}
-exec {giant_exe} analyze compute-one --id "$1" --chunk "$2" --run-dir {run_dir}
-"""
-
-
-def _submit_description(cfg: SubmitConfig, wrapper: Path, jobs_file: Path) -> str:
-    reqs_attrs = (
-        "+RemoteJob = True\n"
-        if cfg.remote
-        else "requirements = TARGET.ProvidesETPResources\n"
-    )
-    return (
-        "universe = docker\n"
-        f"docker_image = {cfg.docker_image}\n"
-        f"executable = {wrapper}\n"
-        "arguments = $(plotid) $(chunk)\n"
-        "should_transfer_files = YES\n"
-        "when_to_transfer_output = ON_EXIT\n"
-        f"request_memory = {cfg.request_memory_mb}\n"
-        f"request_cpus = {cfg.request_cpus}\n"
-        "+RequestWalltime = $(walltime)\n"
-        f"accounting_group = {cfg.accounting_group}\n"
-        f"{reqs_attrs}"
-        f"output = {cfg.run_dir}/logs/$(plotid)__$(chunk).out\n"
-        f"error = {cfg.run_dir}/logs/$(plotid)__$(chunk).err\n"
-        f"log = {cfg.run_dir}/logs/condor.log\n"
-        f"queue plotid,chunk,walltime from {jobs_file}\n"
-    )
-
-
-def _job_walltimes(
-    run_dir: Path, ids: list[str], n_chunks: int
-) -> list[tuple[str, int, int]]:
-    """``(spec_id, chunk, walltime_s)`` for every job, sized from ``run_meta.json``.
-
-    Row counts come from ``prep``'s ``RunMeta.rows_per_chunk``/``total_rows``;
-    ``chunkable=False`` specs (router diagnostics) always use the dataset
-    total since they run as a single job regardless of ``n_chunks``.
-    """
-    meta = RunMeta.load(run_dir / "run_meta.json")
-    jobs: list[tuple[str, int, int]] = []
-    for spec_id in ids:
-        chunkable = get_spec(spec_id).chunkable
-        chunks = range(n_chunks) if chunkable else [0]
-        for chunk in chunks:
-            n_rows = meta.rows_per_chunk[chunk] if chunkable else meta.total_rows
-            jobs.append((spec_id, chunk, estimate_runtime_s(spec_id, n_rows)))
-    return jobs
-
-
-def _resolve_giant_executable(repo_dir: Path) -> Path:
-    """Path to the ``giant`` entry point to bake into the condor wrapper script.
-
-    Prefers the venv currently running this process (``sys.executable``'s
-    sibling ``giant``) so a submit from a non-default venv (e.g. ``--extra
-    cuda`` on a dev box) doesn't silently pick up a different one; falls back
-    to ``repo_dir/.venv/bin/giant`` for the case this is invoked from outside
-    any venv (e.g. a system Python).
-    """
-    active = Path(sys.executable).parent / "giant"
-    if active.exists():
-        return active
-    venv_giant = repo_dir / ".venv" / "bin" / "giant"
-    if not venv_giant.exists():
-        raise FileNotFoundError(
-            f"no `giant` executable found next to {sys.executable} or at "
-            f"{venv_giant} — condor jobs run it directly (no `uv` on the "
-            f"worker image), so run `uv sync --extra cpu` in {repo_dir} "
-            "before submitting."
-        )
-    return venv_giant
-
-
-def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path:
-    """Write the wrapper script, (plot, chunk) job list, and HTCondor submit
-    description.
-
-    Each catalog id gets ``cfg.n_chunks`` jobs, except ``chunkable=False``
-    specs (the router diagnostics), which always get exactly one regardless of
-    ``cfg.n_chunks``. Every job's ``+RequestWalltime`` is estimated from its
-    chunk's row count (``runtime_estimate.estimate_runtime_s``, requires
-    ``run_meta.json`` from ``prep`` to already carry ``rows_per_chunk``).
-    Returns the submit description path (``<run_dir>/analyze.sub``). Does not
-    submit — call ``condor_submit`` on the returned file.
-
-    ``cfg.n_chunks`` and the run directory's own ``RunMeta.n_chunks`` (fixed by
-    ``prep``, and what ``RunMeta.rows_per_chunk`` was sized against) are two
-    independent values — checked equal up front so a mismatch is a clear error
-    here rather than an ``IndexError`` out of ``_job_walltimes``.
-    """
-    giant_exe = _resolve_giant_executable(cfg.repo_dir)
-
-    ids = ids or catalog_ids()
-    run_dir = cfg.run_dir
-    meta = RunMeta.load(run_dir / "run_meta.json")
-    if cfg.n_chunks != meta.n_chunks:
-        raise ValueError(
-            f"SubmitConfig.n_chunks={cfg.n_chunks} does not match the "
-            f"n_chunks this run directory was prepped with "
-            f"(RunMeta.n_chunks={meta.n_chunks} in {run_dir}/run_meta.json) — "
-            "re-run `prep` with the desired n_chunks, or fix cfg.n_chunks to "
-            "match it."
-        )
-    (run_dir / "logs").mkdir(parents=True, exist_ok=True)
-    (run_dir / "reduced").mkdir(parents=True, exist_ok=True)
-    (run_dir / "reduced_partial").mkdir(parents=True, exist_ok=True)
-
-    wrapper = run_dir / "run_compute.sh"
-    wrapper.write_text(
-        _WRAPPER.format(repo_dir=cfg.repo_dir, giant_exe=giant_exe, run_dir=run_dir)
-    )
-    wrapper.chmod(0o755)
-
-    jobs = _job_walltimes(run_dir, ids, cfg.n_chunks)
-    jobs_file = run_dir / "jobs.txt"
-    jobs_file.write_text("\n".join(f"{i},{k},{w}" for i, k, w in jobs) + "\n")
-
-    sub = run_dir / "analyze.sub"
-    sub.write_text(_submit_description(cfg, wrapper, jobs_file))
-    return sub
@@ -1,4 +1,4 @@
-"""Per-(plot, chunk) HTCondor walltime estimates for `giant analyze submit`.
+"""Per-(plot, chunk) HTCondor walltime estimates for the analysis compute jobs.

 Each catalog spec's compute cost is close to linear in the number of input
 rows a `compute-one` job streams over — every spec is one (or a couple of)
@@ -6,8 +6,8 @@ streaming `group_by` pass(es) over the chunk (see `catalog.py`/`reduce.py`).
 `_COST_MODEL` below is ``spec_id -> (intercept_s, seconds_per_row)``.
 ``n_rows`` is the combined rollout+reference row count of the job's input:
 the chunk's row count for `chunkable=True` specs, the whole dataset's for the
-three `chunkable=False` router specs (they always run as a single job
-regardless of chunk count).
+`chunkable=False` router specs in `_ROUTER_IDS` (they always run as a single
+job regardless of chunk count).

 Calibrated 2026-07-27 from real HTCondor timings (`condor_history`
 ``RemoteWallClockTime``) of a production run: prediction ``563f5ee3``
@@ -27,7 +27,7 @@ would then wrongly scale up with a bigger dataset. `RUNTIME_SAFETY_MARGIN` is
 deliberately generous (4x total) specifically to absorb that kind of
 contention spike instead. Rerun this calibration (pull fresh
 `condor_history`/`run_meta.json`, refit) if the catalog changes or timings
-drift — a synthetic local rebaseline via `scripts/profile_analysis_costs.py`
+drift — a synthetic local rebaseline via `giant/tools/profile_analysis_costs.py`
 is a reasonable fallback when no real cluster data is available yet, but
 undershoots real wall time badly (it can't see docker pull / `/ceph` I/O
 latency), which is exactly why this file moved off it.
@@ -54,9 +54,7 @@ _FIXED_OVERHEAD_S = 60.0
 # scan. Calibrated from the 3 real router jobs' observed wall times (119, 66,
 # 124s) — max minus _FIXED_OVERHEAD_S, on top of it.
 _ROUTER_FIXED_S = 64.0
-_ROUTER_IDS = frozenset(
-    {"router_gating", "router_share_by_pdg", "router_share_by_process"}
-)
+_ROUTER_IDS = frozenset({"router_gating", "router_share_by_pdg", "router_share_by_process", "router_specialization"})

 # Conservative fallback for any catalog id not in _COST_MODEL (e.g. a plot
 # added after the last calibration run) — the most expensive fitted per-row
@@ -1,7 +1,12 @@
-"""Canonical world-frame LazyFrame builders for the two sides of a comparison.
+"""Canonical world-frame LazyFrame builders for the two kinds of comparison input.

-The analysis compares one autoregressive ``giant rollout`` (the *generated* side)
-against a raw miniCaloSim steps file (the *reference* / real side). Both carry a
+The analysis compares one or more autoregressive ``giant rollout`` runs (the
+*generated* side — one named series each, see ``RolloutSpec``) against a single
+raw miniCaloSim steps file shared by all of them (the *reference* / real side).
+Every rollout is the same *kind* of file regardless of how many there are, so
+``Side`` stays binary: it describes a file's schema (rollout column layout +
+synthetic-termination rows + per-track secondary view, vs. reference
+``sec_*_list`` columns), not series identity. Both kinds carry a
 **shared world-frame physical column subset** under identical names, so no
 renaming or coordinate decode is needed — everything is already in world-frame
 mm / MeV:
@@ -26,6 +31,7 @@ HTCondor workers that have no LaTeX toolchain.

 from __future__ import annotations

+from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path

@@ -40,6 +46,11 @@ from giant.constants import (
    TERM_MAX_STEPS,
    TERM_UNKNOWN_PDG,
 )
+from giant.data.loader import event_id_offset, find_parquet_files
+
+# Helper column name for the per-shard offset join in open_side; dropped before
+# the LazyFrame is returned, so it never leaks into a caller's schema.
+_SOURCE_PATH_COL = "__source_path"

 # The world-frame physical columns both sides share under identical names.
 PHYS_COLS: tuple[str, ...] = (
@@ -77,12 +88,44 @@ SYNTHETIC_TERMINATION_REASONS: frozenset[str] = frozenset(


 class Side(str, Enum):
-    """Which of the two comparison inputs a file is."""
+    """Which of the two comparison-input *kinds* a file is."""

    rollout = "rollout"
    reference = "reference"


+@dataclass
+class RolloutSpec:
+    """One named rollout input, as fed to ``build_context``/``Bundle.open``.
+
+    ``name`` is the series' identity throughout the rest of the pipeline (a
+    plot's ``payload["series"]`` key, a figure's legend label, its color) —
+    resolved once in ``condor.load_rollout_yamls`` from ``--label`` or the
+    YAML stem, then threaded through unchanged. ``checkpoint`` /
+    ``type_embedding_l1_dist`` are only used by the router/type-embedding
+    diagnostics (``catalog.py``'s ``chunkable=False`` specs).
+    """
+
+    name: str
+    source: str | Path | pl.LazyFrame
+    checkpoint: str | None = None
+    type_embedding_l1_dist: dict | None = None
+
+
+@dataclass
+class RolloutSide:
+    """One rollout's opened frames + per-checkpoint diagnostic inputs (``catalog.Bundle.rollouts`` value)."""
+
+    all: pl.LazyFrame  # rollout, all rows (incl. synthetic termination rows)
+    phys: pl.LazyFrame  # rollout, physical steps only
+    checkpoint: str | None = None  # from the rollout YAML; router_gating only
+    # Diagnostic pre-aggregated at rollout time (giant.rollout.
+    # L1DistCollector.summary()) — from the rollout YAML, type_embedding_l1_distance
+    # only. Unlike checkpoint, this needs no live model: it's already a
+    # finished histogram, just passed through.
+    type_embedding_l1_dist: dict | None = None
+
+
 def _check_rollout_metadata(path: Path) -> None:
    """Raise if ``path`` carries coord metadata that isn't the rollout tag.

@@ -93,10 +136,7 @@ def _check_rollout_metadata(path: Path) -> None:
    metadata = pq.read_schema(path).metadata or {}
    coord = metadata.get(PREDICT_COORD_METADATA_KEY.encode())
    if coord is not None and coord.decode() != ROLLOUT_COORD_VALUE:
-        raise ValueError(
-            f"{path} is not a rollout file (coord={coord.decode()!r}); "
-            "expected `giant rollout` output"
-        )
+        raise ValueError(f"{path} is not a rollout file (coord={coord.decode()!r}); expected `giant rollout` output")


 def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
@@ -111,6 +151,22 @@ def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
    reference file's upstream ROOT→parquet conversion don't agree on integer
    width, and an uncast mismatch only surfaces later as a ``pl.concat``
    ``SchemaError`` (e.g. in ``build_context``'s pdg-count merge).
+
+    The reference (a rollout's seed ``dataset``) may be a directory of parquet
+    shards, or a ``.manifest`` naming a subset, rather than a single file — each
+    such shard is a separate Geant4 job whose own ``event_id`` numbering
+    restarts from 0, so a multi-shard load offsets every shard's ids by
+    ``giant.data.loader.event_id_offset(file_index)`` to keep them globally
+    unique, exactly as the training/rollout data pipeline already does
+    (``giant/data/loader.py``). ``file_index`` comes from
+    ``find_parquet_files``'s deterministic ordering — the same list and
+    ordering ``giant rollout`` used (via ``_seed_from_data``) to offset the
+    rollout side's own ``event_id``s, so both sides agree on what an
+    ``event_id`` means. There is no overflow guard here (unlike
+    ``loader._offset_event_id``): checking it would cost an eager
+    ``event_id``-column read per shard in every condor compute job, and
+    ``giant rollout`` already ran that check over this exact file list when it
+    produced the seed.
    """
    if isinstance(source, pl.LazyFrame):
        return source.with_columns(pl.col("pdg").cast(pl.Int64))
@@ -119,13 +175,18 @@ def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
        _check_rollout_metadata(path)
        lf = pl.scan_parquet(path)
    else:
-        # The reference (a rollout's seed `dataset`) may be a directory of
-        # parquet shards rather than a single file — scan them all.
-        lf = (
-            pl.scan_parquet(str(path / "**/*.parquet"))
-            if path.is_dir()
-            else pl.scan_parquet(path)
-        )
+        files = find_parquet_files(path)
+        if len(files) == 1:
+            lf = pl.scan_parquet(files[0])
+        else:
+            offsets = {str(p): event_id_offset(i) for i, p in enumerate(files)}
+            lf = (
+                pl.scan_parquet(files, include_file_paths=_SOURCE_PATH_COL)
+                .with_columns(
+                    pl.col("event_id") + pl.col(_SOURCE_PATH_COL).replace_strict(offsets, return_dtype=pl.Int64)
+                )
+                .drop(_SOURCE_PATH_COL)
+            )
    return lf.with_columns(pl.col("pdg").cast(pl.Int64))


@@ -137,9 +198,7 @@ def physical_steps(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
    """
    if side is Side.reference:
        return lf
-    return lf.filter(
-        ~pl.col("termination_reason").is_in(list(SYNTHETIC_TERMINATION_REASONS))
-    )
+    return lf.filter(~pl.col("termination_reason").is_in(list(SYNTHETIC_TERMINATION_REASONS)))


 def secondaries(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
@@ -1,4 +1,4 @@
-"""Secondary-type embedding-distance diagnostic (docs/v0.3.0-design.md §11.3).
+"""Secondary-type embedding-distance diagnostic.

 Unlike every other diagnostic in this package, the data isn't derivable from
 a rollout/reference parquet at all — it's the L1 distance between each
@@ -20,26 +20,37 @@ redesign exists to fix.

 from __future__ import annotations

+from typing import TYPE_CHECKING
+
 from giant.analysis.reduced import Reduced

+if TYPE_CHECKING:
+    from giant.analysis.sources import RolloutSide
+
 _NOTE_NOT_APPLICABLE = (
-    "not applicable: this rollout's checkpoint doesn't use "
+    "not applicable: none of these rollouts' checkpoints use "
    "stage2_model.particle_type.target='embedding' (or generated no "
    "secondaries), so giant rollout recorded no type_embedding_l1_dist "
-    "diagnostic in its YAML sidecar"
+    "diagnostic in their YAML sidecar"
 )


-def compute_type_embedding_l1_distance(l1_dist: dict | None) -> Reduced:
-    """`Reduced` for the type-embedding-distance figure, or an explanatory
-    note if this checkpoint never populated the diagnostic.
+def compute_type_embedding_l1_distance(rollouts: dict[str, "RolloutSide"]) -> Reduced:
+    """`Reduced` for the type-embedding-distance figure: one series per rollout
+    whose checkpoint populated the diagnostic, or an explanatory note if none did.

-    `l1_dist`: `giant.rollout.L1DistCollector.summary()`'s dict, as recorded
-    in the rollout YAML's `type_embedding_l1_dist` key (`Bundle.
-    type_embedding_l1_dist`) — `{"n", "mean", "std", "min", "max",
-    "hist_edges", "hist_counts"}`.
+    Each rollout's `RolloutSide.type_embedding_l1_dist` is
+    `giant.rollout.L1DistCollector.summary()`'s dict, as recorded in that
+    rollout's YAML `type_embedding_l1_dist` key — `{"n", "mean", "std",
+    "min", "max", "hist_edges", "hist_counts"}`. Every collector uses the
+    same fixed log-spaced edges (`L1DistCollector.__init__`'s defaults, never
+    overridden — see `giant/cli.py`'s rollout command), so it's safe to plot
+    every rollout's counts against the first one's edges.
    """
-    if l1_dist is None:
+    entries = {
+        name: rs.type_embedding_l1_dist for name, rs in rollouts.items() if rs.type_embedding_l1_dist is not None
+    }
+    if not entries:
        return Reduced(
            id="type_embedding_l1_distance",
            family="model",
@@ -49,6 +60,11 @@ def compute_type_embedding_l1_distance(l1_dist: dict | None) -> Reduced:
            payload={"note": _NOTE_NOT_APPLICABLE},
        )

+    edges = next(iter(entries.values()))["hist_edges"]
+    notes = [
+        f"{name}: n={d['n']:,} mean={d['mean']:.4g} std={d['std']:.4g} min={d['min']:.4g} max={d['max']:.4g}"
+        for name, d in entries.items()
+    ]
    return Reduced(
        id="type_embedding_l1_distance",
        family="model",
@@ -56,15 +72,10 @@ def compute_type_embedding_l1_distance(l1_dist: dict | None) -> Reduced:
        title="Secondary-type embedding L1 distance (predicted vector -> nearest PDG row)",
        xlabel="L1 distance",
        payload={
-            "edges": l1_dist["hist_edges"],
-            "rollout": l1_dist["hist_counts"],
+            "edges": edges,
+            "series": {name: d["hist_counts"] for name, d in entries.items()},
            "log_y": True,
            "log_x": True,
-            "note": (
-                f"n={l1_dist['n']:,}  mean={l1_dist['mean']:.4g}  "
-                f"std={l1_dist['std']:.4g}  min={l1_dist['min']:.4g}  "
-                f"max={l1_dist['max']:.4g}; rollout only, no reference "
-                "concept for a raw pre-decode vector"
-            ),
+            "note": "; ".join(notes) + "; rollout only, no reference concept for a raw pre-decode vector",
        },
    )
@@ -0,0 +1,234 @@
+"""Load a trained checkpoint into ready-to-run models (giant.cli's `predict`/`rollout`).
+
+Both commands need the same ~15 steps to go from a checkpoint path to two
+`eval()`-mode models plus their normalizers/vocab maps: load the pickle,
+validate it carries what current code expects, resolve which conditioning
+mode each axis was trained with, restore the top-N vocab maps (if the
+checkpoint used one-hot conditioning), rebuild the normalizers, construct the
+model from `model_config`, and load the requested (raw or EMA) weights. This
+used to be duplicated near-verbatim in both commands (issues.md Issue 5) —
+`load_for_inference` is the single implementation.
+
+This module intentionally has no Typer dependency, so it can be unit-tested
+directly and imported from non-CLI code (`giant.analysis.router_gating`,
+lazily — see that module's docstring for why). Failures raise
+`CheckpointCompatibilityError` with the same wording the CLI has always
+shown; the CLI layer catches it and does the `typer.echo`/`Exit(1)`.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+from torch import nn
+
+from giant import config as gconfig
+from giant.constants import K_MAX
+from giant.data.loader import TopNMap
+from giant.data.setup_cache import topnmap_from_json
+from giant.data.transforms import Normalizer
+from giant.model.network import build_models
+
+
+class CheckpointCompatibilityError(Exception):
+    """Checkpoint is missing something `load_for_inference` needs."""
+
+
+def conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]:
+    """(particle_conditioning, material_conditioning) for
+    `giant.data.transforms.build_cond_features`/`build_features` — from
+    either a v0.2 checkpoint's flat `model_config["conditioning"]` (one
+    shared string, same for both axes) or a new-format one (independent
+    `model_config["conditioning"]["particle"/"material"]["type"]` — the two
+    axes are configured independently and may differ)."""
+    raw = model_cfg.get("conditioning", default)
+    if isinstance(raw, dict):
+        return (
+            raw.get("particle", {}).get("type", default),
+            raw.get("material", {}).get("type", default),
+        )
+    return raw, raw
+
+
+def stage_cfg(model_cfg: dict, stage: str) -> dict:
+    """`model_cfg[f"{stage}_model"]` for a new-format model_config, `{}` for
+    a v0.2 flat one (whose ddpm schedule always used `CosineSchedule`'s own
+    default `T=1000` — never a config key — and which never had
+    `particle_type` at all, so `{}` is the correct fallback for both
+    `ddpm_steps`/`particle_type_other_policy` below)."""
+    val = model_cfg.get(f"{stage}_model")
+    return val if isinstance(val, dict) else {}
+
+
+def ddpm_steps(model_cfg: dict, stage: str) -> int:
+    return stage_cfg(model_cfg, stage).get("ddpm", {}).get("n_steps", 1000)
+
+
+def particle_type_other_policy(model_cfg: dict) -> str:
+    return stage_cfg(model_cfg, "stage2").get("particle_type", {}).get("other_policy", "sample")
+
+
+def load_pdg_topn_map(ckpt: dict) -> TopNMap | None:
+    """`ckpt["pdg_topn_map"]` as a `giant.data.loader.TopNMap`, or `None` if
+    this checkpoint's conditioning/particle_type never needed one (see
+    `giant.pipeline.run_setup_stage`, which only populates it when
+    `conditioning.particle.type` or `stage2_model.particle_type.target` is
+    `"onehot"`)."""
+    raw = ckpt.get("pdg_topn_map")
+    return topnmap_from_json(raw, axis="pdg") if raw is not None else None
+
+
+def load_mat_topn_map(ckpt: dict) -> TopNMap | None:
+    """`ckpt["mat_topn_map"]` as a `giant.data.loader.TopNMap`, or `None` if
+    this checkpoint's `conditioning.material.type` was never `"onehot"` (see
+    `giant.pipeline.run_setup_stage`)."""
+    raw = ckpt.get("mat_topn_map")
+    return topnmap_from_json(raw, axis="material") if raw is not None else None
+
+
+def load_sec_type_topn_map(ckpt: dict) -> TopNMap | None:
+    """`ckpt["sec_type_topn_map"]` as a `giant.data.loader.TopNMap`, or
+    `None` if this checkpoint's `stage2_model.particle_type.target` was never
+    `"onehot"` (see `giant.pipeline.run_setup_stage`).
+
+    Pre-gitea-#29 checkpoints have no `sec_type_topn_map` key at all — before
+    #29, the secondary-species decode map and the conditioning PDG onehot map
+    were always numerically the same map, saved once under `pdg_topn_map`.
+    For those, fall back to `load_pdg_topn_map` to reproduce that exact
+    behavior; a current checkpoint always has the key (possibly `null`, if
+    `particle_type.target != "onehot"`), so this fallback never fires for one."""
+    if "sec_type_topn_map" in ckpt:
+        raw = ckpt["sec_type_topn_map"]
+        return topnmap_from_json(raw, axis="pdg") if raw is not None else None
+    return load_pdg_topn_map(ckpt)
+
+
+@dataclass(frozen=True)
+class InferenceContext:
+    """Everything needed to run a trained checkpoint forward, resolved once."""
+
+    stage1: nn.Module | None
+    stage2: nn.Module | None
+    cond_norm: Normalizer
+    tgt_norm: Normalizer
+    sec_phys_norm: Normalizer
+    pdg_map: dict[int, int]
+    mat_map: dict[str, int]
+    pdg_topn_map: TopNMap | None
+    mat_topn_map: TopNMap | None
+    sec_type_topn_map: TopNMap | None
+    particle_conditioning: str
+    material_conditioning: str
+    k_max: int
+    stage1_ddpm_steps: int
+    stage2_ddpm_steps: int
+    other_policy: str
+    model_config: dict
+    epoch: int | None
+    best_val_loss: float | None
+
+
+def load_for_inference(
+    checkpoint: Path,
+    device: torch.device,
+    command_name: str,
+    weights: str = "raw",
+    require_stage2: bool = True,
+) -> InferenceContext:
+    """Load *checkpoint* and reconstruct everything `predict`/`rollout` need
+    to run it forward, on *device*, in `eval()` mode.
+
+    *command_name* (e.g. `"predict"`/`"rollout"`) only feeds the "needs both"
+    error message below. *weights* is `"raw"` (the live training weights) or
+    `"ema"` (the EMA shadow copy, see `--ema-decay`). *require_stage2*
+    controls whether a checkpoint with an inactive stage 2
+    (`stage2_model.active = false`) is an error (both current callers need
+    both stages) or an acceptable `stage2 = None` result — kept as a real
+    parameter since `stage{1,2}_model.active` is a real, if currently
+    stage1+stage2-only-in-practice, config option.
+    """
+    ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
+    for key in ("model_config", "sec_decoder"):
+        if key not in ckpt:
+            raise CheckpointCompatibilityError(f"checkpoint has no {key} — retrain with the current code")
+
+    if "sec_phys" not in ckpt.get("normalizer", {}):
+        raise CheckpointCompatibilityError("checkpoint has no normalizer.sec_phys — retrain with the current code")
+
+    gconfig.warn_if_checkpoint_config_mismatch(checkpoint)
+
+    model_cfg = ckpt["model_config"]
+    particle_conditioning, material_conditioning = conditioning_axes(model_cfg)
+    pdg_topn_map = load_pdg_topn_map(ckpt)
+    mat_topn_map = load_mat_topn_map(ckpt)
+    if particle_conditioning == "onehot" and pdg_topn_map is None:
+        raise CheckpointCompatibilityError(
+            "checkpoint's conditioning.particle.type='onehot' but has no pdg_topn_map — retrain with the current code"
+        )
+    if material_conditioning == "onehot" and mat_topn_map is None:
+        raise CheckpointCompatibilityError(
+            "checkpoint's conditioning.material.type='onehot' but has no mat_topn_map — retrain with the current code"
+        )
+    sec_type_topn_map = load_sec_type_topn_map(ckpt)
+    particle_type_target = stage_cfg(model_cfg, "stage2").get("particle_type", {}).get("target", "onehot")
+    if particle_type_target == "onehot" and sec_type_topn_map is None:
+        raise CheckpointCompatibilityError(
+            "checkpoint's stage2_model.particle_type.target='onehot' but has no "
+            "sec_type_topn_map — retrain with the current code"
+        )
+    other_policy = particle_type_other_policy(model_cfg)
+    stage1_ddpm_steps = ddpm_steps(model_cfg, "stage1")
+    stage2_ddpm_steps = ddpm_steps(model_cfg, "stage2")
+    k_max = stage_cfg(model_cfg, "stage2").get("k_max", K_MAX)
+    pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
+    mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
+    cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
+    tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
+    sec_phys_norm = Normalizer.from_dict(ckpt["normalizer"]["sec_phys"])
+
+    built = build_models(model_cfg)
+    stage1, stage2 = built["stage1"], built["stage2"]
+    if require_stage2 and (stage1 is None or stage2 is None):
+        raise CheckpointCompatibilityError(
+            f"checkpoint has an inactive stage1 or stage2 — {command_name} needs both (see stage{{1,2}}_model.active)"
+        )
+
+    if weights == "raw":
+        model_key, sec_key = "model", "sec_decoder"
+    else:
+        model_key, sec_key = "model_ema", "sec_decoder_ema"
+        if model_key not in ckpt or sec_key not in ckpt:
+            raise CheckpointCompatibilityError(
+                f"{checkpoint} has no EMA weights (trained before --ema-decay, "
+                "or with --ema-decay 0) — use --weights raw"
+            )
+    if stage1 is not None:
+        stage1.load_state_dict(ckpt[model_key])
+        stage1.to(device).eval()
+    if stage2 is not None:
+        stage2.load_state_dict(ckpt[sec_key])
+        stage2.to(device).eval()
+
+    return InferenceContext(
+        stage1=stage1,
+        stage2=stage2,
+        cond_norm=cond_norm,
+        tgt_norm=tgt_norm,
+        sec_phys_norm=sec_phys_norm,
+        pdg_map=pdg_map,
+        mat_map=mat_map,
+        pdg_topn_map=pdg_topn_map,
+        mat_topn_map=mat_topn_map,
+        sec_type_topn_map=sec_type_topn_map,
+        particle_conditioning=particle_conditioning,
+        material_conditioning=material_conditioning,
+        k_max=k_max,
+        stage1_ddpm_steps=stage1_ddpm_steps,
+        stage2_ddpm_steps=stage2_ddpm_steps,
+        other_policy=other_policy,
+        model_config=model_cfg,
+        epoch=ckpt.get("epoch"),
+        best_val_loss=ckpt.get("best_val_loss"),
+    )
@@ -0,0 +1,103 @@
+"""Single source of truth for the conditioning arrays' column layout (gitea #37).
+
+`cond_cont` and `cond_cat` are built in `giant.data.transforms` and consumed in
+`giant.model.encoders` / `giant.model.routers`. Their column order used to be
+written down independently on each side, kept in sync only by parallel comments
+— so getting it wrong produced silently mis-indexed columns rather than an
+exception, and adding a conditioning axis meant a coordinated multi-file edit.
+
+`CondLayout` owns that order. Both sides construct one from the same
+`conditioning.particle.type` / `conditioning.material.type` pair and read named
+slices off it, so the layout is stated exactly once. This module depends only on
+`giant.constants`, so both the data and model packages can import it.
+"""
+
+from dataclasses import dataclass
+from typing import ClassVar
+
+from giant.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM
+
+# The three per-axis conditioning modes. Mirrors giant.config.Conditioning,
+# which this module deliberately does not import (giant.config pulls in the
+# whole model package).
+AXIS_TYPES = ("physical", "embedding", "onehot")
+
+
+@dataclass(frozen=True)
+class CondLayout:
+    """Column layout of `cond_cont`/`cond_cat` for one (particle, material) mode pair.
+
+    `cond_cont` is unconditionally `COND_DIM` wide regardless of mode: the base
+    block, then the particle physical block, then the material physical block.
+    An axis that isn't `"physical"` gets its block zero-filled and never reads
+    it (see `giant.data.transforms._physical_cond_columns`), so the widths are
+    mode-independent and only the *meaning* of a block changes.
+
+    `cond_cat` is 2 to 4 wide. Columns `PDG_COL`/`MAT_COL` are always the dense
+    training-vocab index; an axis in `"onehot"` mode appends one more column
+    holding its top-N-plus-other class index, particle before material.
+    """
+
+    particle_type: str
+    material_type: str
+
+    # cond_cat's dense-vocab columns, present in every mode. Under
+    # "physical"/"onehot" they are a reporting/router convenience the
+    # ConditionEncoder never reads; under "embedding" they are the signal.
+    PDG_COL: ClassVar[int] = 0
+    MAT_COL: ClassVar[int] = 1
+
+    def __post_init__(self) -> None:
+        if self.particle_type not in AXIS_TYPES:
+            raise ValueError(f"unknown conditioning.particle.type {self.particle_type!r}")
+        if self.material_type not in AXIS_TYPES:
+            raise ValueError(f"unknown conditioning.material.type {self.material_type!r}")
+
+    @classmethod
+    def from_types(cls, particle_type: str, material_type: str) -> "CondLayout":
+        """Named constructor — the entry point both sides use."""
+        return cls(particle_type=particle_type, material_type=material_type)
+
+    # --- cond_cont ---------------------------------------------------------
+
+    @property
+    def base(self) -> slice:
+        """pre_pos(3), log(pre_E)(1), pre_dir(3), layer_id(1)."""
+        return slice(0, COND_DIM_BASE)
+
+    @property
+    def particle_phys(self) -> slice:
+        """log(mass), charge — see `giant.particles`."""
+        return slice(COND_DIM_BASE, COND_DIM_BASE + PARTICLE_PHYS_DIM)
+
+    @property
+    def material_phys(self) -> slice:
+        """Z_eff, A_eff, log(density), log(X0), log(lambda_int) — see `giant.materials`."""
+        start = COND_DIM_BASE + PARTICLE_PHYS_DIM
+        return slice(start, start + MATERIAL_PHYS_DIM)
+
+    @property
+    def cont_dim(self) -> int:
+        return COND_DIM
+
+    # --- cond_cat ----------------------------------------------------------
+
+    @property
+    def particle_topn_col(self) -> int | None:
+        """Column of the particle top-N class index, or `None` if not `"onehot"`."""
+        return self.MAT_COL + 1 if self.particle_type == "onehot" else None
+
+    @property
+    def material_topn_col(self) -> int | None:
+        """Column of the material top-N class index, or `None` if not `"onehot"`.
+
+        Comes after the particle top-N column when both axes are `"onehot"`.
+        """
+        if self.material_type != "onehot":
+            return None
+        return self.MAT_COL + (2 if self.particle_type == "onehot" else 1)
+
+    @property
+    def cat_dim(self) -> int:
+        """Total `cond_cat` width: 2, plus one column per `"onehot"` axis."""
+        return self.MAT_COL + 1 + (self.particle_type == "onehot") + (self.material_type == "onehot")
@@ -1,6 +1,7 @@
 from __future__ import annotations

 from pathlib import Path
+from typing import NamedTuple

 import numpy as np
 import torch
@@ -11,6 +12,37 @@ from giant.data.loader import event_id_offset, iter_file_chunks
 from giant.data.transforms import Normalizer, build_features, sorted_membership


+class StepBatch(NamedTuple):
+    """One training batch, as yielded by `StreamingStepsDataset`. Field order
+    is load-bearing for existing positional unpacking elsewhere (`trainers.py`,
+    `validate.py`, test fixtures) — append only, never insert or reorder.
+
+    cond_cont:  (B, COND_DIM) float32
+    cond_cat:   (B, 2/3/4) int64 — width 2 unless conditioning="onehot"
+    target_s1:  (B, 9) float32 — normalised Stage-1 primary target
+    n_sec:      (B,) int64 — true secondary count per step
+    sec_cont:   (B, k_max, SEC_SLOT_DIM) float32 — [stick_logit,
+                local_dir, log_mass, charge] per slot (mass/charge
+                normalised iff `sec_phys_normalizer` was given); always
+                computed the same way regardless of
+                stage2_model.particle_type.target, only actually used
+                downstream under target="physical"
+    proc_idx:   (B,) int64 — process-class label (ProcessRouter supervision
+                only; zeros when `proc_map` is None)
+    sec_type_idx: (B, k_max) int64 — per-slot class index into
+                `sec_type_class_map`, for particle_type.target in
+                ("onehot", "embedding"); zeros (unused) otherwise
+    """
+
+    cond_cont: torch.Tensor
+    cond_cat: torch.Tensor
+    target_s1: torch.Tensor
+    n_sec: torch.Tensor
+    sec_cont: torch.Tensor
+    proc_idx: torch.Tensor
+    sec_type_idx: torch.Tensor
+
+
 def make_event_split(
    all_event_ids: np.ndarray,
    val_fraction: float = 0.1,
@@ -39,28 +71,10 @@ class StreamingStepsDataset(IterableDataset):
    rather than single rows, so the batch is assembled with vectorized
    numpy slicing instead of a per-row Python loop in the default collate.

-    Each batch is a tuple:
-        (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, sec_type_idx)
-    where:
-        cond_cont:  (B, COND_DIM) float32
-        cond_cat:   (B, 2/3/4) int64 — width 2 unless conditioning="onehot"
-                    (see docs/v0.3.0-design.md decision 4)
-        target_s1:  (B, 9) float32 — normalised Stage-1 primary target
-        n_sec:      (B,) int64 — true secondary count per step
-        sec_cont:   (B, k_max, SEC_SLOT_DIM) float32 — [stick_logit,
-                    local_dir, log_mass, charge] per slot (mass/charge
-                    normalised iff `sec_phys_normalizer` was given); always
-                    computed the same way regardless of
-                    stage2_model.particle_type.target, only actually used
-                    downstream under target="physical"
-        proc_idx:   (B,) int64 — process-class label (ProcessRouter supervision
-                    only; zeros when `proc_map` is None)
-        sec_type_idx: (B, k_max) int64 — per-slot class index into
-                    `sec_type_class_map`, for particle_type.target in
-                    ("onehot", "embedding"); zeros (unused) otherwise
+    Each batch is a `StepBatch` — see its docstring for field meanings.

    `k_max` (constructor arg, default the module constant) should match
-    `stage2_model.k_max` (docs/v0.3.0-design.md §8) — it sets the padded
+    `stage2_model.k_max` — it sets the padded
    width of `sec_cont`/`sec_type_idx` above.
    """

@@ -83,6 +97,7 @@ class StreamingStepsDataset(IterableDataset):
        mat_topn_map: dict[str, int] | None = None,
        sec_type_class_map: dict | None = None,
        k_max: int = K_MAX,
+        seed: int = 0,
    ) -> None:
        self.files = list(files)
        self._offsets = {path: event_id_offset(i) for i, path in enumerate(self.files)}
@@ -103,16 +118,35 @@ class StreamingStepsDataset(IterableDataset):
        self.mat_topn_map = mat_topn_map
        self.sec_type_class_map = sec_type_class_map
        self.k_max = k_max
+        self.seed = seed
+        self.epoch = 0
+        self._rng = np.random.default_rng()
+
+    def set_epoch(self, epoch: int) -> None:
+        """Select the shuffle stream for `epoch` (the DistributedSampler convention).
+
+        The training loop calls this at the top of every epoch. Shuffling is
+        seeded from `(seed, epoch, worker_id)` rather than the global numpy
+        state so epoch *k*'s batch order is the same whether it runs as epoch
+        *k* of one long `giant train`, or as its own resumed job in a
+        per-epoch workflow chain (`giant/workflow/tasks.py:TrainEpochTask`).
+        Workers are re-forked from this object each epoch (no
+        `persistent_workers`), so setting it here reaches them.
+        """
+        self.epoch = int(epoch)

    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        files = self.files
+        worker_id = worker_info.id if worker_info is not None else 0
        if worker_info is not None:
-            files = files[worker_info.id :: worker_info.num_workers]
+            files = files[worker_id :: worker_info.num_workers]
+
+        self._rng = np.random.default_rng([self.seed, self.epoch, worker_id])

        if self.shuffle:
            files = list(files)
-            np.random.default_rng().shuffle(files)
+            self._rng.shuffle(files)

        buf_cont: list[np.ndarray] = []
        buf_cat: list[np.ndarray] = []
@@ -124,25 +158,13 @@ class StreamingStepsDataset(IterableDataset):
        buf_n = 0

        for path in files:
-            for chunk in iter_file_chunks(
-                path, offset=self._offsets[path], k_max=self.k_max
-            ):
+            for chunk in iter_file_chunks(path, offset=self._offsets[path], k_max=self.k_max):
                mask = sorted_membership(chunk["event_id"], self._events_arr)
                if not mask.any():
                    continue
                chunk = {k: v[mask] for k, v in chunk.items()}

-                (
-                    cond_cont,
-                    cond_cat,
-                    target_s1,
-                    n_sec,
-                    sec_cont,
-                    proc_idx,
-                    sec_type_idx,
-                    _,
-                    _,
-                ) = build_features(
+                feats = build_features(
                    chunk,
                    self.pdg_map,
                    self.mat_map,
@@ -158,14 +180,14 @@ class StreamingStepsDataset(IterableDataset):
                    sec_type_class_map=self.sec_type_class_map,
                    k_max=self.k_max,
                )
-                buf_cont.append(cond_cont)
-                buf_cat.append(cond_cat)
-                buf_tgt.append(target_s1)
-                buf_nsec.append(n_sec)
-                buf_sec.append(sec_cont)
-                buf_proc.append(proc_idx)
-                buf_type.append(sec_type_idx)
-                buf_n += len(cond_cont)
+                buf_cont.append(feats.cond_cont)
+                buf_cat.append(feats.cond_cat)
+                buf_tgt.append(feats.target_s1)
+                buf_nsec.append(feats.n_sec)
+                buf_sec.append(feats.sec_cont)
+                buf_proc.append(feats.proc_idx)
+                buf_type.append(feats.sec_type_idx)
+                buf_n += len(feats.cond_cont)

                if buf_n >= self.shuffle_buffer:
                    (
@@ -220,7 +242,7 @@ class StreamingStepsDataset(IterableDataset):
        styp = np.concatenate(buf_type)

        if self.shuffle:
-            idx = np.random.permutation(len(cont))
+            idx = self._rng.permutation(len(cont))
            cont, cat, tgt = cont[idx], cat[idx], tgt[idx]
            nsec, sec, proc, styp = nsec[idx], sec[idx], proc[idx], styp[idx]

@@ -229,14 +251,14 @@ class StreamingStepsDataset(IterableDataset):
        n_full = n // bs if not final else (n + bs - 1) // bs
        for start in range(0, n_full * bs, bs):
            end = min(start + bs, n)
-            yield (
-                torch.from_numpy(cont[start:end]).float(),
-                torch.from_numpy(cat[start:end]).long(),
-                torch.from_numpy(tgt[start:end]).float(),
-                torch.from_numpy(nsec[start:end]).long(),
-                torch.from_numpy(sec[start:end]).float(),
-                torch.from_numpy(proc[start:end]).long(),
-                torch.from_numpy(styp[start:end]).long(),
+            yield StepBatch(
+                cond_cont=torch.from_numpy(cont[start:end]).float(),
+                cond_cat=torch.from_numpy(cat[start:end]).long(),
+                target_s1=torch.from_numpy(tgt[start:end]).float(),
+                n_sec=torch.from_numpy(nsec[start:end]).long(),
+                sec_cont=torch.from_numpy(sec[start:end]).float(),
+                proc_idx=torch.from_numpy(proc[start:end]).long(),
+                sec_type_idx=torch.from_numpy(styp[start:end]).long(),
            )

        if final:
@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Iterator

@@ -16,7 +16,7 @@ from giant.constants import K_MAX
 MANIFEST_SUFFIX = ".manifest"

 # Each input parquet file is a separate Geant4 job converted 1:1 from its own
-# ROOT file (scripts/steps_to_parquet.py), and a job's event_id numbering
+# ROOT file (giant/tools/steps_to_parquet.py), and a job's event_id numbering
 # always restarts from 0 — so when multiple files are loaded together (a
 # directory or .manifest), raw event_id values collide across files even
 # though they refer to unrelated events. Every per-file event_id column gets
@@ -115,9 +115,7 @@ def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndar
    return out


-def _df_to_dict(
-    df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX
-) -> dict[str, np.ndarray]:
+def _df_to_dict(df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
    has_sec_lists = "sec_E_list" in df.columns

    d: dict[str, np.ndarray] = {
@@ -136,9 +134,7 @@ def _df_to_dict(
        # / ProcessRouter). Guarded like has_sec_lists: older parquet
        # conversions predating this column still load fine.
        "process": (
-            df["process"].to_numpy(dtype=object)
-            if "process" in df.columns
-            else np.full(len(df), "", dtype=object)
+            df["process"].to_numpy(dtype=object) if "process" in df.columns else np.full(len(df), "", dtype=object)
        ),
        "step_length": df["step_length"].to_numpy(dtype=np.float32),
        "post_E": df["post_E"].to_numpy(dtype=np.float32),
@@ -151,16 +147,12 @@ def _df_to_dict(
    if has_sec_lists:
        d["sec_E_list"] = _pad_list_col(df["sec_E_list"], k_max)
        d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], k_max)
-        d["sec_dir_list"] = _pad_dir_col(
-            df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], k_max
-        )
+        d["sec_dir_list"] = _pad_dir_col(df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], k_max)

    return d


-def load_steps(
-    path: str | Path, offset: int = 0, k_max: int = K_MAX
-) -> dict[str, np.ndarray]:
+def load_steps(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
    return _df_to_dict(pd.read_parquet(path), offset=offset, k_max=k_max)


@@ -170,13 +162,11 @@ def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
    return _offset_event_id(ids, offset)


-def iter_file_chunks(
-    path: str | Path, offset: int = 0, k_max: int = K_MAX
-) -> Iterator[dict[str, np.ndarray]]:
+def iter_file_chunks(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> Iterator[dict[str, np.ndarray]]:
    """Yield one parquet row-group at a time so a large file never fully loads.

    `k_max` sets the padded width of the sec_*_list columns (should match
-    `stage2_model.k_max` — see docs/v0.3.0-design.md §8); defaults to the
+    `stage2_model.k_max`); defaults to the
    module constant for callers that don't care (e.g. Stage-1-only reads)."""
    pf = pq.ParquetFile(path)
    for i in range(pf.num_row_groups):
@@ -214,15 +204,11 @@ def _cond_df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]
    }


-def iter_cond_chunks(
-    path: str | Path, offset: int = 0
-) -> Iterator[dict[str, np.ndarray]]:
+def iter_cond_chunks(path: str | Path, offset: int = 0) -> Iterator[dict[str, np.ndarray]]:
    """Yield conditioning-only row-groups (no post-step columns read from disk)."""
    pf = pq.ParquetFile(path)
    for i in range(pf.num_row_groups):
-        yield _cond_df_to_dict(
-            pf.read_row_group(i, columns=_COND_COLS).to_pandas(), offset=offset
-        )
+        yield _cond_df_to_dict(pf.read_row_group(i, columns=_COND_COLS).to_pandas(), offset=offset)


 def build_index_maps(
@@ -270,25 +256,32 @@ def _rank_by_frequency_from_files(files: list[Path], column: str, cast) -> dict:
    return counts


-def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict]:
+def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict, dict]:
    """Frequency-capped value->index map: the `n_classes - 1` most frequent
    keys get their own index; every rarer key is bucketed into a shared
    "other" index (`n_classes - 1`).

-    Returns `(class_map, other_members)` — `other_members` is `{key: count}`
-    for every key bucketed into "other" (the empirical within-bucket
-    distribution, for `other_policy = "sample"` at rollout — see
-    docs/v0.3.0-design.md §8).
+    Returns `(class_map, other_members, class_counts)` — `other_members` is
+    `{key: count}` for every key bucketed into "other" (the empirical
+    within-bucket distribution, for `other_policy = "sample"` at rollout);
+    `class_counts` is `{index: total_count}` for every resulting class index
+    (0-indexed; the "other" index's count is the sum of `other_members`),
+    the per-class frequencies `stage2_model.particle_type.class_weighting`
+    (gitea #44) needs and that would otherwise be dropped once `counts` is
+    collapsed into `class_map`.
    """
    ranked = sorted(counts, key=lambda k: counts[k], reverse=True)
    keep = ranked[: max(n_classes - 1, 0)]
    class_map = {k: i for i, k in enumerate(keep)}
+    class_counts = {i: counts[k] for i, k in enumerate(keep)}
    other_idx = n_classes - 1
    other_members: dict = {}
    for k in ranked[len(keep) :]:
        class_map[k] = other_idx
        other_members[k] = counts[k]
-    return class_map, other_members
+    if other_members:
+        class_counts[other_idx] = sum(other_members.values())
+    return class_map, other_members, class_counts


 def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str, int]:
@@ -302,7 +295,7 @@ def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str,
    fixed-width n_sec_head classifier.
    """
    counts = _rank_by_frequency_from_files(files, "process", str)
-    class_map, _ = _topn_plus_other_map(counts, n_experts)
+    class_map, _, _ = _topn_plus_other_map(counts, n_experts)
    return class_map


@@ -314,17 +307,20 @@ class TopNMap:

    class_map: dict
    other_members: dict
+    # {class_index: total_count} — see _topn_plus_other_map. Empty for a
+    # TopNMap decoded from a checkpoint/sidecar predating gitea #44; only
+    # stage2_model.particle_type.class_weighting reads it, and it raises
+    # loudly if it needs counts that aren't there (giant/training/trainers.py).
+    class_counts: dict = field(default_factory=dict)


-def build_topn_map_from_files(
-    files: list[Path], column: str, n_classes: int, cast=str
-) -> TopNMap:
+def build_topn_map_from_files(files: list[Path], column: str, n_classes: int, cast=str) -> TopNMap:
    """Scan `column` and build a frequency-capped value->index map, structurally
    identical to `build_process_map_from_files` (shares its ranking core via
    `_topn_plus_other_map`), generalized over the source column and key type.

    Used for the material axis (`column="material"`, `cast=str`, matching
-    `mat_map`'s key type) — see docs/v0.3.0-design.md §8. The PDG axis uses
+    `mat_map`'s key type). The PDG axis uses
    `build_pdg_topn_map_from_files` instead (it needs to pool two columns,
    which this single-column form can't express). Also records
    `other_members` (the empirical within-"other" distribution), needed
@@ -332,8 +328,8 @@ def build_topn_map_from_files(
    free during this same scan.
    """
    counts = _rank_by_frequency_from_files(files, column, cast)
-    class_map, other_members = _topn_plus_other_map(counts, n_classes)
-    return TopNMap(class_map=class_map, other_members=other_members)
+    class_map, other_members, class_counts = _topn_plus_other_map(counts, n_classes)
+    return TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)


 def build_pdg_topn_map_from_files(files: list[Path], n_classes: int) -> TopNMap:
@@ -341,8 +337,8 @@ def build_pdg_topn_map_from_files(files: list[Path], n_classes: int) -> TopNMap:
    plays in this dataset: a step's own primary particle (`pdg` column) and
    an emitted secondary's species (`sec_pdg_list`, exploded) — shared by
    `conditioning.particle.type = "onehot"` and
-    `stage2_model.particle_type.target = "onehot"` (docs/v0.3.0-design.md
-    §8). Pooling both is what keeps a species that's common as a secondary
+    `stage2_model.particle_type.target = "onehot"`. Pooling both is what
+    keeps a species that's common as a secondary
    but rare as a primary (or vice versa) from being pushed into "other"
    just because one role's count alone looks small — the meeting's failure
    mode (zero photon secondaries, hallucinated antineutrinos) was
@@ -364,5 +360,5 @@ def build_pdg_topn_map_from_files(files: list[Path], n_classes: int) -> TopNMap:
        if has_sec:
            exploded = df["sec_pdg_list"].explode().dropna()
            _accumulate_value_counts(counts, exploded, int)
-    class_map, other_members = _topn_plus_other_map(counts, n_classes)
-    return TopNMap(class_map=class_map, other_members=other_members)
+    class_map, other_members, class_counts = _topn_plus_other_map(counts, n_classes)
+    return TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)
@@ -35,7 +35,10 @@ from giant.data.transforms import Normalizer, sorted_membership
 # v3: NormalizerEntry.energy_reservoir_sample (100k raw values) replaced by
 # energy_quantiles (a fixed ENERGY_QUANTILE_LEVELS-point quantile grid) — a
 # v2 sidecar has no such grid to fall back on, so it must be recomputed.
-_CACHE_FORMAT_VERSION = 3
+# v4: TopNMap gained class_counts (gitea #44, stage2_model.particle_type.
+# class_weighting) — a v3 sidecar's cached topn_maps have no counts, so they
+# must be rebuilt rather than silently cached with class_counts={}.
+_CACHE_FORMAT_VERSION = 4

 _DIMS = {
    "COND_DIM": COND_DIM,
@@ -104,17 +107,14 @@ def normalizer_key(
 ) -> str:
    # .6g avoids float-repr drift (e.g. 0.1 vs 0.10000000000000002) causing
    # spurious cache misses between runs with the "same" val_fraction. The two
-    # conditioning axes are independent (docs/v0.3.0-design.md §3.1) and both
+    # conditioning axes are independent and both
    # affect which cond_cont columns are computed for real vs. zero-filled
    # (giant.data.transforms._physical_cond_columns), so both must be part of
    # the key or two mixed-axis runs could collide on the same cache entry.
-    return (
-        f"valfrac={val_fraction:.6g}_seed={seed}_pcond={particle_conditioning}"
-        f"_mcond={material_conditioning}"
-    )
+    return f"valfrac={val_fraction:.6g}_seed={seed}_pcond={particle_conditioning}_mcond={material_conditioning}"


-# Top-N-map axes (docs/v0.3.0-design.md §8): "pdg" keys match pdg_map's int
+# Top-N-map axes: "pdg" keys match pdg_map's int
 # keys (shared by conditioning.particle.type="onehot" and
 # stage2_model.particle_type.target="onehot" — one map for both), "material"
 # keys match mat_map's str keys.
@@ -126,10 +126,7 @@ def topn_key(axis: str, n_classes: int) -> str:
    sidecar stays reusable across runs with different emb_dim (see the
    dict[int, dict] precedent `proc_maps` sets, keyed by n_experts)."""
    if axis not in _TOPN_AXIS_CASTS:
-        raise ValueError(
-            f"unknown top-N map axis {axis!r}, expected one of "
-            f"{sorted(_TOPN_AXIS_CASTS)}"
-        )
+        raise ValueError(f"unknown top-N map axis {axis!r}, expected one of {sorted(_TOPN_AXIS_CASTS)}")
    return f"{axis}:{n_classes}"


@@ -137,6 +134,7 @@ def topnmap_to_json(m: TopNMap) -> dict:
    return {
        "class_map": {str(k): v for k, v in m.class_map.items()},
        "other_members": {str(k): v for k, v in m.other_members.items()},
+        "class_counts": {str(k): v for k, v in m.class_counts.items()},
    }


@@ -145,6 +143,11 @@ def topnmap_from_json(d: dict, axis: str) -> TopNMap:
    return TopNMap(
        class_map={cast(k): v for k, v in d["class_map"].items()},
        other_members={cast(k): v for k, v in d["other_members"].items()},
+        # Missing for a checkpoint's topn maps predating gitea #44 — {} is
+        # the correct decode there (inference never reads class_counts; only
+        # stage2_model.particle_type.class_weighting does, at train time, and
+        # it raises loudly if it needs counts a checkpoint doesn't have).
+        class_counts={int(k): v for k, v in d.get("class_counts", {}).items()},
    )


@@ -165,9 +168,7 @@ class NormalizerEntry:
            "tgt_norm": self.tgt_norm.to_dict(),
            "sec_phys_norm": self.sec_phys_norm.to_dict(),
            "n_train_steps": self.n_train_steps,
-            "energy_quantiles": np.asarray(
-                self.energy_quantiles, dtype=np.float32
-            ).tolist(),
+            "energy_quantiles": np.asarray(self.energy_quantiles, dtype=np.float32).tolist(),
        }

    @classmethod
@@ -190,7 +191,7 @@ class SetupCache:
    proc_maps: dict[int, dict[str, int]] = field(default_factory=dict)
    normalizers: dict[str, NormalizerEntry] = field(default_factory=dict)
    topn_maps: dict[str, TopNMap] = field(default_factory=dict)
-    """Keyed by `topn_key(axis, n_classes)` — see docs/v0.3.0-design.md §8."""
+    """Keyed by `topn_key(axis, n_classes)`."""

    @classmethod
    def empty(cls, files: list[Path]) -> "SetupCache":
@@ -234,13 +235,8 @@ class SetupCache:
                np.array(d["event_index"]["counts"], dtype=np.int64),
            )
        proc_maps = {int(k): v for k, v in d.get("proc_maps", {}).items()}
-        normalizers = {
-            k: NormalizerEntry.from_json(v) for k, v in d.get("normalizers", {}).items()
-        }
-        topn_maps = {
-            k: topnmap_from_json(v, axis=k.split(":", 1)[0])
-            for k, v in d.get("topn_maps", {}).items()
-        }
+        normalizers = {k: NormalizerEntry.from_json(v) for k, v in d.get("normalizers", {}).items()}
+        topn_maps = {k: topnmap_from_json(v, axis=k.split(":", 1)[0]) for k, v in d.get("topn_maps", {}).items()}
        return cls(
            fingerprint=d["fingerprint"],
            git_hash=d.get("git_hash", "unknown"),
@@ -263,18 +259,14 @@ class SetupCache:
            fingerprint=other.fingerprint,
            git_hash=other.git_hash,
            vocab=other.vocab if other.vocab is not None else self.vocab,
-            event_index=(
-                other.event_index if other.event_index is not None else self.event_index
-            ),
+            event_index=(other.event_index if other.event_index is not None else self.event_index),
            proc_maps={**self.proc_maps, **other.proc_maps},
            normalizers={**self.normalizers, **other.normalizers},
            topn_maps={**self.topn_maps, **other.topn_maps},
        )


-def load(
-    data: str | Path, files: list[Path], echo=lambda *a, **k: None
-) -> SetupCache | None:
+def load(data: str | Path, files: list[Path], echo=lambda *a, **k: None) -> SetupCache | None:
    """Load and validate the sidecar for `data`; `None` on any miss (never raises).

    A missing file, corrupt JSON, format-version mismatch, dimension-constant
@@ -298,9 +290,7 @@ def load(
            echo("setup cache: format version changed — ignoring stale cache")
            return None
        if raw.get("dims") != _DIMS:
-            echo(
-                "setup cache: model dimension constants changed — ignoring stale cache"
-            )
+            echo("setup cache: model dimension constants changed — ignoring stale cache")
            return None
        fp = fingerprint_files(files)
        if raw.get("fingerprint") != fp:
@@ -343,9 +333,7 @@ def save(
        with open(lock_path, "a") as lock_file:
            fcntl.flock(lock_file, fcntl.LOCK_EX)
            try:
-                base = load(data, files, echo=lambda *a, **k: None) or SetupCache.empty(
-                    files
-                )
+                base = load(data, files, echo=lambda *a, **k: None) or SetupCache.empty(files)
                merged = base.merge(sections)
                payload = json.dumps(merged.to_json(), separators=(",", ":"))
                tmp.write_text(payload)
@@ -353,9 +341,7 @@ def save(
            finally:
                fcntl.flock(lock_file, fcntl.LOCK_UN)
    except OSError as exc:
-        echo(
-            f"setup cache: could not write {path} ({exc}) — continuing without caching"
-        )
+        echo(f"setup cache: could not write {path} ({exc}) — continuing without caching")
        try:
            tmp.unlink(missing_ok=True)
        except OSError:
@@ -366,16 +352,12 @@ def compute_event_index_from_files(files: list[Path]) -> tuple[np.ndarray, np.nd
    """Unique event ids + per-event row (step) counts, across all `files`."""
    if not files:
        return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64)
-    all_ids = np.concatenate(
-        [load_event_ids(f, offset=event_id_offset(i)) for i, f in enumerate(files)]
-    )
+    all_ids = np.concatenate([load_event_ids(f, offset=event_id_offset(i)) for i, f in enumerate(files)])
    unique_ids, counts = np.unique(all_ids, return_counts=True)
    return unique_ids, counts


-def n_train_steps_for_split(
-    unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray
-) -> int:
+def n_train_steps_for_split(unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray) -> int:
    """Row (step) count summed over whichever `unique_ids` fall in `train_events_arr`.

    `train_events_arr` must be ascending and duplicate-free (as produced by
@@ -1,7 +1,9 @@
 import warnings
+from typing import NamedTuple

 import numpy as np

+from giant.cond_layout import CondLayout
 from giant.constants import K_MAX

 _EPS = 1e-8
@@ -12,6 +14,14 @@ _EPS = 1e-8
 # the conservation it slightly softens is physically negligible (~0.001%).
 _SIMPLEX_FLOOR = 1e-5

+# Upper clip for a raw predicted log_mass before inv_log_transform: exp(y)
+# must stay well inside float32 range (~3.4e38, i.e. y < ~88.7) or it
+# overflows to inf, which — like the negative-mass case below — blows up the
+# next log_transform call once that mass is fed back in as conditioning.
+# 80.0 leaves comfortable headroom while still being far beyond any physical
+# particle mass a converged model would ever predict.
+_LOG_MASS_MAX = 80.0
+

 def log_transform(x: np.ndarray, eps: float = _EPS) -> np.ndarray:
    x = np.asarray(x, dtype=np.float32)
@@ -84,9 +94,7 @@ def energy_simplex_encode(
    return z.astype(np.float32)


-def energy_simplex_decode(
-    z: np.ndarray, pre_E: np.ndarray
-) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+def energy_simplex_decode(z: np.ndarray, pre_E: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Inverse of `energy_simplex_encode`: ALR coords + pre_E → physical energies.

    A softmax over `[z_edep, z_sec, 0]` recovers the three simplex fractions, so
@@ -118,9 +126,7 @@ def _rodrigues_axis(pre_dir: np.ndarray) -> np.ndarray:
    arbitrary second operand) on every row; profiling on a 114M-row file
    showed `np.cross` as the single hottest call inside this rotation.
    """
-    axis = np.stack(
-        [pre_dir[:, 1], -pre_dir[:, 0], np.zeros_like(pre_dir[:, 0])], axis=1
-    )
+    axis = np.stack([pre_dir[:, 1], -pre_dir[:, 0], np.zeros_like(pre_dir[:, 0])], axis=1)
    axis_norm = np.linalg.norm(axis, axis=1, keepdims=True)  # (N,1)
    # axis_norm ~ 0 happens at BOTH poles: pre_dir ~ +ẑ (forward) and
    # pre_dir ~ -ẑ (near-exact backscatter) — ‖pre_dir × ẑ‖ = sin(angle to
@@ -199,9 +205,7 @@ def local_frame_rotation(pre_dir: np.ndarray, post_dir: np.ndarray) -> np.ndarra
    kxv = _cross_with_z_axis(axis, post_dir)  # (N,3)
    kdv = (axis * post_dir).sum(axis=1, keepdims=True)  # (N,1)

-    return (post_dir * cos_t + kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(
-        np.float32
-    )
+    return (post_dir * cos_t + kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(np.float32)


 class Normalizer:
@@ -341,9 +345,7 @@ def sorted_membership(values: np.ndarray, sorted_arr: np.ndarray) -> np.ndarray:
    return sorted_arr[idx] == values


-def _vectorized_map_lookup(
-    values: np.ndarray, mapping: dict, strict: bool = True, default: int = 0
-) -> np.ndarray:
+def _vectorized_map_lookup(values: np.ndarray, mapping: dict, strict: bool = True, default: int = 0) -> np.ndarray:
    """Vectorized equivalent of `np.array([mapping[v] for v in values], dtype=np.int64)`.

    Replaces a per-element Python dict lookup with one `searchsorted` call.
@@ -387,9 +389,7 @@ def travel_direction(pre_pos: np.ndarray, post_pos: np.ndarray) -> np.ndarray:
    disp = post_pos - pre_pos
    norm = np.linalg.norm(disp, axis=1, keepdims=True)
    safe_norm = np.where(norm < 1e-7, 1.0, norm)
-    return np.where(norm < 1e-7, np.array([[0.0, 0.0, 1.0]]), disp / safe_norm).astype(
-        np.float32
-    )
+    return np.where(norm < 1e-7, np.array([[0.0, 0.0, 1.0]]), disp / safe_norm).astype(np.float32)


 def reconstruct_post_pos(
@@ -408,9 +408,7 @@ def reconstruct_post_pos(
    return (pre_pos + step_length.reshape(-1, 1) * travel_dir_world).astype(np.float32)


-def inv_local_frame_rotation(
-    pre_dir: np.ndarray, post_dir_local: np.ndarray
-) -> np.ndarray:
+def inv_local_frame_rotation(pre_dir: np.ndarray, post_dir_local: np.ndarray) -> np.ndarray:
    """Inverse of local_frame_rotation: rotate from local frame back to world frame.

    Applies R^T (same axis, negative angle) to post_dir_local.
@@ -424,9 +422,7 @@ def inv_local_frame_rotation(
    kdv = (axis * post_dir_local).sum(axis=1, keepdims=True)

    # Negative angle: sin_t → -sin_t
-    return (post_dir_local * cos_t - kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(
-        np.float32
-    )
+    return (post_dir_local * cos_t - kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(np.float32)


 _STICK_LOGIT_CLIP = 10.0  # logit value used for the last valid secondary slot
@@ -491,14 +487,10 @@ def encode_secondaries(
                remaining_raw = e_sec - cumsum[:, i - 1]
            shortfall_flagged |= sec_valid[:, i] & (remaining_raw < -_SHORTFALL_TOL)
            remaining = np.maximum(remaining_raw, _EPS)
-            f = np.clip(
-                sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS
-            )
+            f = np.clip(sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS)
            logit = np.log(f / (1.0 - f)).astype(np.float32)
            # Last valid slot: give it the full remaining budget
-            is_last = sec_valid[:, i] & ~(
-                sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool)
-            )
+            is_last = sec_valid[:, i] & ~(sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool))
            logit = np.where(is_last, _STICK_LOGIT_CLIP, logit)
            logit = np.where(
                sec_valid[:, i],
@@ -525,9 +517,7 @@ def encode_secondaries(
            # Only rotate valid slots; leave padded slots as (0,0,1) or whatever.
            valid_mask = sec_valid[:, i]
            if valid_mask.any():
-                dir_local[valid_mask, i] = local_frame_rotation(
-                    pre_dir[valid_mask], sec_dir_list[valid_mask, i]
-                )
+                dir_local[valid_mask, i] = local_frame_rotation(pre_dir[valid_mask], sec_dir_list[valid_mask, i])

    if sec_pdg_list is not None:
        from giant.particles import particle_phys_array
@@ -553,17 +543,15 @@ def encode_secondaries(
    return sec_cont.astype(np.float32)


-def encode_secondary_type_idx(
-    sec_pdg_list: np.ndarray, sec_valid: np.ndarray, class_map: dict
-) -> np.ndarray:
+def encode_secondary_type_idx(sec_pdg_list: np.ndarray, sec_valid: np.ndarray, class_map: dict) -> np.ndarray:
    """Per-secondary-slot class index into `class_map` — (N, K_MAX) int64.

    `class_map` is either a top-N-plus-other map's `class_map`
    (`stage2_model.particle_type.target = "onehot"`, see
    `giant.data.loader.build_pdg_topn_map_from_files`) or the dense `pdg_map`
-    (`target = "embedding"`). Not used at all for `target = "physical"`
-    (see docs/v0.3.0-design.md decision 1) — that target keeps using
-    `encode_secondaries`'s (log_mass, charge) columns unchanged.
+    (`target = "embedding"`). Not used at all for `target = "physical"` —
+    that target keeps using `encode_secondaries`'s (log_mass, charge)
+    columns unchanged.

    Padding slots get index 0 (their looked-up value is discarded downstream
    by the `sec_valid`/`n_sec` mask regardless, so any in-vocabulary dummy
@@ -583,9 +571,7 @@ def encode_secondary_type_idx(
    # isn't guaranteed to be a key — an arbitrary present one always is).
    dummy = next(iter(class_map))
    safe_pdg = np.where(sec_valid, sec_pdg_list, dummy)
-    idx = _vectorized_map_lookup(safe_pdg.reshape(-1), class_map, strict=True).reshape(
-        N, K
-    )
+    idx = _vectorized_map_lookup(safe_pdg.reshape(-1), class_map, strict=True).reshape(N, K)
    return np.where(sec_valid, idx, 0).astype(np.int64)


@@ -599,7 +585,7 @@ def decode_secondary_cont(
    stick-breaking energy split and local->world direction — generator/
    `particle_type.target`-independent, since every target (`"physical"`,
    `"onehot"`, `"embedding"`) shares the same `CONT_SLOT_DIM`-wide
-    (stick_logit, dir) prefix (docs/v0.3.0-design.md §6.1) and differs only
+    (stick_logit, dir) prefix and differs only
    in what follows it. `decode_secondaries` (target="physical") is the
    original all-in-one form built on top of this; `target` in `("onehot",
    "embedding")` decodes their type slice separately via
@@ -659,9 +645,7 @@ def decode_secondary_cont(
    for i in range(K):
        valid = sec_valid[:, i]
        if valid.any():
-            sec_dir_world[valid, i] = inv_local_frame_rotation(
-                pre_dir[valid], dir_local[valid, i]
-            )
+            sec_dir_world[valid, i] = inv_local_frame_rotation(pre_dir[valid], dir_local[valid, i])

    return sec_E, sec_dir_world, sec_valid

@@ -704,35 +688,36 @@ def decode_secondaries(
        sec_cont = sec_cont.copy()
        sec_cont[:, :, 4:6] = phys.reshape(N_, K_, 2)

-    sec_E, sec_dir_world, sec_valid = decode_secondary_cont(
-        sec_cont, n_sec, e_sec, pre_dir
-    )
+    sec_E, sec_dir_world, sec_valid = decode_secondary_cont(sec_cont, n_sec, e_sec, pre_dir)

    log_mass = sec_cont[:, :, 4]  # (N, K)
    charge = sec_cont[:, :, 5]  # (N, K)

-    # mass is non-negative by construction (inv_log_transform of a real
-    # number is always > 0); clip to 0 for padded/invalid slots rather than
-    # leaving a spurious small positive floor from the log inverse.
+    # log_mass is a raw model prediction, not itself the output of
+    # log_transform, so it can land far outside the range that round-trips
+    # cleanly through inv_log_transform: too negative and exp(log_mass)
+    # undershoots _EPS, making inv_log_transform go slightly negative; too
+    # positive and exp(log_mass) overflows float32 to inf. Either one then
+    # blows up the next log_transform call on this track's mass once it's
+    # fed back in as conditioning for a further rollout step
+    # (giant/rollout.py -> build_cond_features -> _physical_cond_columns).
+    # Clip to a range whose inverse is guaranteed finite and >= 0 before
+    # that can happen; clip to 0 separately for padded/invalid slots rather
+    # than leaving a spurious small positive floor.
+    log_mass = np.clip(log_mass, np.log(_EPS), _LOG_MASS_MAX)
    sec_mass = np.where(sec_valid, inv_log_transform(log_mass), 0.0).astype(np.float32)
    sec_charge = np.where(sec_valid, charge, 0.0).astype(np.float32)

    return sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid


-def _physical_cond_columns(
-    data: dict[str, np.ndarray],
-    particle_conditioning: str,
-    material_conditioning: str,
-) -> np.ndarray:
+def _physical_cond_columns(data: dict[str, np.ndarray], layout: CondLayout) -> np.ndarray:
    """(N, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM) physical conditioning columns.

-    The particle and material blocks are gated independently
-    (docs/v0.3.0-design.md §3.1: "configured independently and may mix
-    freely — e.g. material `physical` with particle `embedding`"), so e.g.
-    `particle_conditioning="embedding"` + `material_conditioning="physical"`
-    zero-fills only the particle columns and computes the material ones for
-    real.
+    The particle and material blocks are gated independently and may mix
+    freely — e.g. material `physical` with particle `embedding` — so e.g.
+    `particle_type="embedding"` + `material_type="physical"` zero-fills only
+    the particle columns and computes the material ones for real.

    "embedding"/"onehot" zero-fill their block (cheap, and ConditionEncoder
    never reads these columns in either mode — so an unfilled
@@ -749,7 +734,7 @@ def _physical_cond_columns(

    n = len(next(iter(data.values())))

-    if particle_conditioning == "physical":
+    if layout.particle_type == "physical":
        from giant.particles import particle_phys_array

        if "mass" in data and "charge" in data:
@@ -758,19 +743,13 @@ def _physical_cond_columns(
        else:
            mass, charge = particle_phys_array(data["pdg"]).T
        particle_cols = np.column_stack([log_transform(mass), charge])
-    elif particle_conditioning in ("embedding", "onehot"):
-        particle_cols = np.zeros((n, PARTICLE_PHYS_DIM), dtype=np.float32)
    else:
-        raise ValueError(
-            f"unknown conditioning.particle.type {particle_conditioning!r}"
-        )
+        particle_cols = np.zeros((n, PARTICLE_PHYS_DIM), dtype=np.float32)

-    if material_conditioning == "physical":
+    if layout.material_type == "physical":
        from giant.materials import material_properties_array

-        z_eff, a_eff, density, x0, lambda_int = material_properties_array(
-            data["material"]
-        ).T
+        z_eff, a_eff, density, x0, lambda_int = material_properties_array(data["material"]).T
        material_cols = np.column_stack(
            [
                z_eff,
@@ -780,16 +759,66 @@ def _physical_cond_columns(
                log_transform(lambda_int),
            ]
        )
-    elif material_conditioning in ("embedding", "onehot"):
-        material_cols = np.zeros((n, MATERIAL_PHYS_DIM), dtype=np.float32)
    else:
-        raise ValueError(
-            f"unknown conditioning.material.type {material_conditioning!r}"
-        )
+        material_cols = np.zeros((n, MATERIAL_PHYS_DIM), dtype=np.float32)

    return np.column_stack([particle_cols, material_cols]).astype(np.float32)


+def _build_cond_arrays(
+    data: dict[str, np.ndarray],
+    pdg_map: dict[int, int],
+    mat_map: dict[str, int],
+    layout: CondLayout,
+    pdg_topn_map: dict[int, int] | None,
+    mat_topn_map: dict[str, int] | None,
+) -> tuple[np.ndarray, np.ndarray]:
+    """The un-normalized `(cond_cont, cond_cat)` pair, in `layout`'s column order.
+
+    Both `build_cond_features` and `build_features` go through here, so the
+    column order — and everything that depends on it — is stated once. See
+    `giant.cond_layout.CondLayout` for the layout itself.
+    """
+    cond_cont = np.column_stack(
+        [
+            data["pre_pos"],
+            log_transform(data["pre_E"]),
+            data["pre_dir"],
+            data["layer_id"].astype(np.float32),
+        ]
+    ).astype(np.float32)  # (N, COND_DIM_BASE=8)
+    cond_cont = np.column_stack([cond_cont, _physical_cond_columns(data, layout)]).astype(
+        np.float32
+    )  # (N, COND_DIM=15)
+
+    # In "physical" mode cond_cat's first two columns are only a
+    # reporting/router convenience — ConditionEncoder never reads them
+    # (giant/model/encoders.py) — so a species/material outside the training
+    # vocab (the whole point of physical-property conditioning) gets a dummy
+    # index instead of raising. In "embedding" mode those columns ARE the
+    # conditioning signal, so an unmapped value must still raise loudly
+    # rather than silently misassign. In "onehot" mode they again go unread
+    # (the topN columns below are the real signal), so they're as permissive
+    # as "physical". Each axis's strictness is independent.
+    pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=layout.particle_type == "embedding")
+    mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=layout.material_type == "embedding")
+    # Which extra columns exist is the layout's call, not "did the caller
+    # happen to pass a map" — that's what used to let the producer and
+    # ConditionEncoder disagree. A map for a non-"onehot" axis is unused.
+    cat_cols = [pdg_idx, mat_idx]
+    if layout.particle_topn_col is not None:
+        if pdg_topn_map is None:
+            raise ValueError("conditioning.particle.type='onehot' needs pdg_topn_map")
+        cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map))
+    if layout.material_topn_col is not None:
+        if mat_topn_map is None:
+            raise ValueError("conditioning.material.type='onehot' needs mat_topn_map")
+        cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map))
+    cond_cat = np.column_stack(cat_cols)  # (N, layout.cat_dim)
+
+    return cond_cont, cond_cat
+
+
 def build_cond_features(
    data: dict[str, np.ndarray],
    pdg_map: dict[int, int],
@@ -802,57 +831,21 @@ def build_cond_features(
 ) -> tuple[np.ndarray, np.ndarray]:
    """Build conditioning arrays only — no target, no post-step variables.

-    `particle_conditioning`/`material_conditioning` are independent
-    (docs/v0.3.0-design.md §3.1) — e.g. `particle_conditioning="embedding"` +
+    `particle_conditioning`/`material_conditioning` are independent —
+    e.g. `particle_conditioning="embedding"` +
    `material_conditioning="physical"` is a valid mix.

    `pdg_topn_map`/`mat_topn_map` (a top-N-plus-other `class_map`, see
-    `giant.data.loader.build_topn_map_from_files`) append extra `cond_cat`
-    columns read by `ConditionEncoder`'s `"onehot"` mode
-    (docs/v0.3.0-design.md decision 4): pdg topN index at column 2 (iff
-    `pdg_topn_map` given), material topN index at column 3 (iff
-    `mat_topn_map` given, after column 2 if both are). Only ever given when
-    the corresponding axis is `"onehot"`; `cond_cat` stays `(N, 2)` otherwise.
+    `giant.data.loader.build_topn_map_from_files`) supply the extra `cond_cat`
+    columns read by `ConditionEncoder`'s `"onehot"` mode, and are required
+    whenever the corresponding axis is `"onehot"`. See
+    `giant.cond_layout.CondLayout` for which columns exist where.
    """
-    cond_cont = np.column_stack(
-        [
-            data["pre_pos"],
-            log_transform(data["pre_E"]),
-            data["pre_dir"],
-            data["layer_id"].astype(np.float32),
-        ]
-    ).astype(np.float32)
-    cond_cont = np.column_stack(
-        [
-            cond_cont,
-            _physical_cond_columns(data, particle_conditioning, material_conditioning),
-        ]
-    ).astype(np.float32)
-
-    # In "physical" mode cond_cat's first two columns are only a
-    # reporting/router convenience — ConditionEncoder never reads them
-    # (giant/model/network.py) — so a species/material outside the training
-    # vocab (the whole point of physical-property conditioning) gets a dummy
-    # index instead of raising. In "embedding" mode those columns ARE the
-    # conditioning signal, so an unmapped value must still raise loudly
-    # rather than silently misassign. In "onehot" mode they again go unread
-    # (the topN columns below are the real signal), so they're as permissive
-    # as "physical". Each axis's strictness is independent.
-    pdg_strict = particle_conditioning == "embedding"
-    mat_strict = material_conditioning == "embedding"
-    pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=pdg_strict)
-    mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=mat_strict)
-    cat_cols = [pdg_idx, mat_idx]
-    if pdg_topn_map is not None:
-        cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map))
-    if mat_topn_map is not None:
-        cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map))
-    cond_cat = np.column_stack(cat_cols)
+    layout = CondLayout.from_types(particle_conditioning, material_conditioning)
+    cond_cont, cond_cat = _build_cond_arrays(data, pdg_map, mat_map, layout, pdg_topn_map, mat_topn_map)

    if cond_normalizer is not None:
-        cond_cont = _cond_normalizer_transform(
-            cond_cont, cond_normalizer, particle_conditioning, material_conditioning
-        )
+        cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, layout)

    return cond_cont, cond_cat

@@ -860,8 +853,7 @@ def build_cond_features(
 def _cond_normalizer_transform(
    cond_cont: np.ndarray,
    cond_normalizer: "Normalizer",
-    particle_conditioning: str,
-    material_conditioning: str,
+    layout: CondLayout,
 ) -> np.ndarray:
    """Apply ``cond_normalizer``, padding a legacy narrower normalizer if needed.

@@ -869,7 +861,7 @@ def _cond_normalizer_transform(
    8->15, ``giant/constants.py``) saved a ``COND_DIM_BASE``-wide (8) cond
    normalizer, fit before ``build_cond_features`` grew the extra physical
    columns. When NEITHER axis is "physical" those columns are never read by
-    ``ConditionEncoder`` (``giant/model/network.py``), so padding the missing
+    ``ConditionEncoder`` (``giant/model/encoders.py``), so padding the missing
    entries with mean=0/std=1 is a safe no-op that keeps such checkpoints
    usable under the current, always-``COND_DIM``-wide contract. If EITHER
    axis is "physical" its columns are load-bearing, so a mismatch there is a
@@ -880,14 +872,14 @@ def _cond_normalizer_transform(
    width = cond_cont.shape[-1]
    if mean.shape[-1] < width:
        physical_load_bearing = "physical" in (
-            particle_conditioning,
-            material_conditioning,
+            layout.particle_type,
+            layout.material_type,
        )
        if physical_load_bearing:
            raise ValueError(
                f"cond normalizer has {mean.shape[-1]} columns, expected "
-                f"{width}, and particle_conditioning={particle_conditioning!r}/"
-                f"material_conditioning={material_conditioning!r} reads the "
+                f"{width}, and particle_conditioning={layout.particle_type!r}/"
+                f"material_conditioning={layout.material_type!r} reads the "
                "physical columns directly — this checkpoint predates "
                "physical-property conditioning and can't be safely padded; "
                "retrain it under the current code."
@@ -898,6 +890,41 @@ def _cond_normalizer_transform(
    return ((cond_cont - mean) / std).astype(np.float32)


+class StepFeatures(NamedTuple):
+    """Output of `build_features`. Field order is load-bearing for existing
+    positional unpacking (tests, `StreamingStepsDataset`) — append only,
+    never insert or reorder.
+
+    target_s1: (N, 9) Stage-1 primary post-step target (unchanged from Phase 1)
+    n_sec: (N,) integer secondary counts (target for n_sec head)
+    sec_cont: (N, K_MAX, SEC_SLOT_DIM=6) continuous secondary targets
+              [stick_logit, dir_local, log_mass, charge] — mass/charge are
+              the secondary's real physical identity (from its ground-truth
+              PDG code), a fixed regression target, not a learned/snapped one.
+              Always computed the same way regardless of
+              `stage2_model.particle_type.target` — only actually used
+              downstream under `target = "physical"`.
+    proc_idx: (N,) integer process-class label (ProcessRouter supervision only —
+              never conditioning). Zeros when `proc_map` is None or the loaded
+              data has no "process" column (e.g. pre-conversion parquet files).
+    sec_type_idx: (N, K_MAX) integer secondary class index into
+              `sec_type_class_map`, for `stage2_model.particle_type.target`
+              in `("onehot", "embedding")` — see `encode_secondary_type_idx`.
+              Zero-filled (and unused) when `sec_type_class_map` is None
+              (i.e. `target = "physical"`).
+    """
+
+    cond_cont: np.ndarray
+    cond_cat: np.ndarray
+    target_s1: np.ndarray
+    n_sec: np.ndarray
+    sec_cont: np.ndarray
+    proc_idx: np.ndarray
+    sec_type_idx: np.ndarray
+    cond_normalizer: Normalizer | None
+    target_normalizer: Normalizer | None
+
+
 def build_features(
    data: dict[str, np.ndarray],
    pdg_map: dict[int, int],
@@ -915,38 +942,10 @@ def build_features(
    mat_topn_map: dict[str, int] | None = None,
    sec_type_class_map: dict | None = None,
    k_max: int = K_MAX,
-) -> tuple[
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    Normalizer | None,
-    Normalizer | None,
-]:
-    """Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx,
-    sec_type_idx) arrays.
-
-    target_s1: (N, 9) Stage-1 primary post-step target (unchanged from Phase 1)
-    n_sec: (N,) integer secondary counts (target for n_sec head)
-    sec_cont: (N, K_MAX, SEC_SLOT_DIM=6) continuous secondary targets
-              [stick_logit, dir_local, log_mass, charge] — mass/charge are
-              the secondary's real physical identity (from its ground-truth
-              PDG code), a fixed regression target, not a learned/snapped one.
-              Always computed the same way regardless of
-              `stage2_model.particle_type.target` (docs/v0.3.0-design.md
-              decision 1/3) — only actually used downstream under `target =
-              "physical"`.
-    proc_idx: (N,) integer process-class label (ProcessRouter supervision only —
-              never conditioning). Zeros when `proc_map` is None or the loaded
-              data has no "process" column (e.g. pre-conversion parquet files).
-    sec_type_idx: (N, K_MAX) integer secondary class index into
-              `sec_type_class_map`, for `stage2_model.particle_type.target`
-              in `("onehot", "embedding")` — see `encode_secondary_type_idx`.
-              Zero-filled (and unused) when `sec_type_class_map` is None
-              (i.e. `target = "physical"`).
+) -> StepFeatures:
+    """Assemble a `StepFeatures` of (cond_cont, cond_cat, target_s1, n_sec,
+    sec_cont, proc_idx, sec_type_idx, cond_normalizer, target_normalizer) —
+    see `StepFeatures` for field meanings.

    require_secondaries: when True, raise if any step has n_sec > 0 but the
        per-secondary list columns are absent (a mis-converted file that would
@@ -958,7 +957,7 @@ def build_features(
        instead) for callers (normalizer fitting) that only read
        `sec_cont[:, :, 4:6]` and would otherwise discard that work.

-    pdg_topn_map/mat_topn_map: appended `cond_cat` columns for
+    pdg_topn_map/mat_topn_map: source of the extra `cond_cat` columns for
        `ConditionEncoder`'s `"onehot"` mode — see `build_cond_features`.

    sec_type_class_map: the map `sec_type_idx` is looked up against — a
@@ -966,7 +965,7 @@ def build_features(
        dense `pdg_map` for `target = "embedding"` (pass `pdg_map` itself).
        `None` for `target = "physical"`.

-    k_max: should match `stage2_model.k_max` (docs/v0.3.0-design.md §8) —
+    k_max: should match `stage2_model.k_max` —
        overridden internally by `data["sec_E_list"]`'s own padded width when
        present (the loader already padded it to some k_max; that width is
        authoritative), so this only actually matters when secondary list
@@ -975,13 +974,9 @@ def build_features(
    """

    post_dir_local = local_frame_rotation(data["pre_dir"], data["post_dir"])
-    travel_dir_local = local_frame_rotation(
-        data["pre_dir"], travel_direction(data["pre_pos"], data["post_pos"])
-    )
+    travel_dir_local = local_frame_rotation(data["pre_dir"], travel_direction(data["pre_pos"], data["post_pos"]))

-    energy_z = energy_simplex_encode(
-        data["edep"], data["e_sec"], data["post_E"], data["pre_E"]
-    )  # (N, 2)
+    energy_z = energy_simplex_encode(data["edep"], data["e_sec"], data["post_E"], data["pre_E"])  # (N, 2)

    target_s1 = np.column_stack(
        [
@@ -993,33 +988,10 @@ def build_features(
    ).astype(np.float32)  # (N, 9)

    # Phase 2: conditioning drops n_sec and log(e_sec)
-    cond_cont = np.column_stack(
-        [
-            data["pre_pos"],
-            log_transform(data["pre_E"]),
-            data["pre_dir"],
-            data["layer_id"].astype(np.float32),
-        ]
-    ).astype(np.float32)  # (N, COND_DIM_BASE=8)
-    cond_cont = np.column_stack(
-        [
-            cond_cont,
-            _physical_cond_columns(data, particle_conditioning, material_conditioning),
-        ]
-    ).astype(np.float32)  # (N, COND_DIM=15)
+    layout = CondLayout.from_types(particle_conditioning, material_conditioning)
+    cond_cont, cond_cat = _build_cond_arrays(data, pdg_map, mat_map, layout, pdg_topn_map, mat_topn_map)

-    pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map)
-    mat_idx = _vectorized_map_lookup(data["material"], mat_map)
-    cat_cols = [pdg_idx, mat_idx]
-    if pdg_topn_map is not None:
-        cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map))
-    if mat_topn_map is not None:
-        cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map))
-    cond_cat = np.column_stack(cat_cols)  # (N, 2/3/4) — see decision 4
-
-    n_sec_raw = data["n_sec"].astype(
-        np.int64
-    )  # (N,) unclamped, for the valid-slot mask
+    n_sec_raw = data["n_sec"].astype(np.int64)  # (N,) unclamped, for the valid-slot mask

    # Secondary continuous targets
    sec_E_list = data.get("sec_E_list")
@@ -1029,7 +1001,7 @@ def build_features(
        # The loader already padded sec_*_list to some k_max (see
        # giant.data.loader.iter_file_chunks); that padded width is
        # authoritative over whatever this call happened to pass in, so the
-        # two can never drift apart (docs/v0.3.0-design.md §8).
+        # two can never drift apart.
        k_max = sec_E_list.shape[1]

    # Clamp the classification label to k_max: the head only has k_max+1
@@ -1084,7 +1056,7 @@ def build_features(
        target_normalizer = Normalizer().fit(target_s1)

    if cond_normalizer is not None:
-        cond_cont = cond_normalizer.transform(cond_cont)
+        cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, layout)
    if target_normalizer is not None:
        target_s1 = target_normalizer.transform(target_s1)
    if sec_phys_normalizer is not None:
@@ -1099,14 +1071,14 @@ def build_features(
    else:
        proc_idx = np.zeros(len(cond_cat), dtype=np.int64)

-    return (
-        cond_cont,
-        cond_cat,
-        target_s1,
-        n_sec,
-        sec_cont,
-        proc_idx,
-        sec_type_idx,
-        cond_normalizer,
-        target_normalizer,
+    return StepFeatures(
+        cond_cont=cond_cont,
+        cond_cat=cond_cat,
+        target_s1=target_s1,
+        n_sec=n_sec,
+        sec_cont=sec_cont,
+        proc_idx=proc_idx,
+        sec_type_idx=sec_type_idx,
+        cond_normalizer=cond_normalizer,
+        target_normalizer=target_normalizer,
    )
@@ -31,10 +31,7 @@ import numpy as np
 import pandas as pd
 import pyarrow.parquet as pq

-_INSTALL_HINT = (
-    "the geometry oracle needs scikit-learn — install it with "
-    "`uv sync --extra cpu --extra geometry`"
-)
+_INSTALL_HINT = "the geometry oracle needs scikit-learn — install it with `uv sync --extra cpu --extra geometry`"


 def _require_sklearn():
@@ -63,9 +60,7 @@ class _SlabLookup:
    layer_ids: np.ndarray  # (n_segments,) int64, layer_id of each segment
    radius_max: float  # largest transverse radius seen in training data

-    def query(
-        self, pos: np.ndarray, margin: float
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    def query(self, pos: np.ndarray, margin: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        other = [i for i in range(3) if i != self.axis]
        z = pos[:, self.axis]
        radius = np.sqrt(pos[:, other[0]] ** 2 + pos[:, other[1]] ** 2)
@@ -75,11 +70,7 @@ class _SlabLookup:
        material = self.materials[idx]
        layer_id = self.layer_ids[idx]

-        escaped = (
-            (z < self.z_edges[0] - margin)
-            | (z > self.z_edges[-1] + margin)
-            | (radius > self.radius_max + margin)
-        )
+        escaped = (z < self.z_edges[0] - margin) | (z > self.z_edges[-1] + margin) | (radius > self.radius_max + margin)
        return material, layer_id, escaped


@@ -300,16 +291,12 @@ def _fit_slab_lookup(
    """
    other = [i for i in range(3) if i != axis]
    z = pos[:, axis].astype(np.float64)
-    radius = np.sqrt(
-        pos[:, other[0]].astype(np.float64) ** 2
-        + pos[:, other[1]].astype(np.float64) ** 2
-    )
+    radius = np.sqrt(pos[:, other[0]].astype(np.float64) ** 2 + pos[:, other[1]].astype(np.float64) ** 2)

    z_min, z_max = float(z.min()), float(z.max())
    if z_min == z_max:
        raise ValueError(
-            "all points share the same depth-axis coordinate — pick a "
-            "different `depth_axis` or use method='knn'/'svm'"
+            "all points share the same depth-axis coordinate — pick a different `depth_axis` or use method='knn'/'svm'"
        )
    edges = np.linspace(z_min, z_max, n_bins + 1)
    bin_idx = np.clip(np.searchsorted(edges, z, side="right") - 1, 0, n_bins - 1)
@@ -344,12 +331,7 @@ def _fit_slab_lookup(
    bin_layer = bin_layer[fill_from]

    # Run-length-encode consecutive bins sharing a label into segments.
-    changed = (
-        np.flatnonzero(
-            (bin_material[1:] != bin_material[:-1]) | (bin_layer[1:] != bin_layer[:-1])
-        )
-        + 1
-    )
+    changed = np.flatnonzero((bin_material[1:] != bin_material[:-1]) | (bin_layer[1:] != bin_layer[:-1])) + 1
    seg_starts = np.concatenate([[0], changed])
    z_edges = np.concatenate([edges[seg_starts], edges[-1:]])
    materials = bin_material[seg_starts]
@@ -460,11 +442,7 @@ def build_geometry_oracle(
    # Escape threshold from the reference point spacing. Sample a subset for the
    # median 2-NN distance (the 1st neighbour of a training point is itself).
    nn = NearestNeighbors(n_neighbors=2).fit(X)
-    probe = (
-        X
-        if len(X) <= 20_000
-        else X[np.random.default_rng(seed).choice(len(X), 20_000, replace=False)]
-    )
+    probe = X if len(X) <= 20_000 else X[np.random.default_rng(seed).choice(len(X), 20_000, replace=False)]
    d2, _ = nn.kneighbors(probe, n_neighbors=2)
    median_nn = float(np.median(d2[:, 1]))
    escape_threshold = escape_factor * median_nn
@@ -64,18 +64,10 @@ MATERIAL_PROPERTIES: dict[str, MaterialProperties] = {
    "G4_CESIUM_IODIDE": MaterialProperties(
        z_eff=54.0, a_eff=129.904539, density=4.51, x0=1.860288, lambda_int=39.305990
    ),
-    "G4_Pb": MaterialProperties(
-        z_eff=82.0, a_eff=207.216962, density=11.35, x0=0.561253, lambda_int=18.247950
-    ),
-    "G4_W": MaterialProperties(
-        z_eff=74.0, a_eff=183.841648, density=19.30, x0=0.350418, lambda_int=10.311580
-    ),
-    "G4_Cu": MaterialProperties(
-        z_eff=29.0, a_eff=63.545648, density=8.96, x0=1.435578, lambda_int=15.587940
-    ),
-    "G4_Fe": MaterialProperties(
-        z_eff=26.0, a_eff=55.845113, density=7.874, x0=1.757493, lambda_int=16.990300
-    ),
+    "G4_Pb": MaterialProperties(z_eff=82.0, a_eff=207.216962, density=11.35, x0=0.561253, lambda_int=18.247950),
+    "G4_W": MaterialProperties(z_eff=74.0, a_eff=183.841648, density=19.30, x0=0.350418, lambda_int=10.311580),
+    "G4_Cu": MaterialProperties(z_eff=29.0, a_eff=63.545648, density=8.96, x0=1.435578, lambda_int=15.587940),
+    "G4_Fe": MaterialProperties(z_eff=26.0, a_eff=55.845113, density=7.874, x0=1.757493, lambda_int=16.990300),
    "G4_BRASS": MaterialProperties(
        z_eff=30.939130,
        a_eff=68.500857,
@@ -83,9 +75,7 @@ MATERIAL_PROPERTIES: dict[str, MaterialProperties] = {
        x0=1.367465,
        lambda_int=16.947420,
    ),
-    "G4_POLYSTYRENE": MaterialProperties(
-        z_eff=3.5, a_eff=6.509339, density=1.06, x0=41.312510, lambda_int=68.749880
-    ),
+    "G4_POLYSTYRENE": MaterialProperties(z_eff=3.5, a_eff=6.509339, density=1.06, x0=41.312510, lambda_int=68.749880),
    "G4_PLASTIC_SC_VINYLTOLUENE": MaterialProperties(
        z_eff=3.368421,
        a_eff=6.219791,
@@ -108,20 +98,15 @@ MATERIAL_PROPERTIES: dict[str, MaterialProperties] = {
        x0=30392.070000,
        lambda_int=71009.500000,
    ),
-    "G4_lAr": MaterialProperties(
-        z_eff=18.0, a_eff=39.947692, density=1.396, x0=14.003440, lambda_int=85.706400
-    ),
+    "G4_lAr": MaterialProperties(z_eff=18.0, a_eff=39.947692, density=1.396, x0=14.003440, lambda_int=85.706400),
 }


-def get_material_properties(
-    name: str, table: dict[str, MaterialProperties] | None = None
-) -> MaterialProperties:
+def get_material_properties(name: str, table: dict[str, MaterialProperties] | None = None) -> MaterialProperties:
    t = MATERIAL_PROPERTIES if table is None else table
    if name not in t:
        raise UnknownMaterialError(
-            f"material {name!r} is not in giant.materials.MATERIAL_PROPERTIES "
-            f"-- add it (known: {sorted(t)})"
+            f"material {name!r} is not in giant.materials.MATERIAL_PROPERTIES -- add it (known: {sorted(t)})"
        )
    props = t[name]
    if any(v is None for v in props):
@@ -134,9 +119,7 @@ def get_material_properties(
    return props


-def material_properties_array(
-    names: np.ndarray, table: dict[str, MaterialProperties] | None = None
-) -> np.ndarray:
+def material_properties_array(names: np.ndarray, table: dict[str, MaterialProperties] | None = None) -> np.ndarray:
    """(N,) str material names -> (N, 5) float32 [z_eff, a_eff, density, x0, lambda_int]."""
    out = np.array(
        [get_material_properties(str(m), table) for m in np.asarray(names)],
@@ -0,0 +1,114 @@
+"""v0.2 -> v0.3 checkpoint migration: translates a v0.2 checkpoint's flat
+`model_config`/state dicts into the current nested shape (issues.md Issue 8;
+see also `giant._migration` and `giant.config.migrate_config`, the sibling
+config.toml migration surface — issues.md Issue 6)."""
+
+from giant._migration import V02_FIXED_FACTS, reject_legacy_router_expert_sizing
+from giant.constants import EMB_DIM, K_MAX
+
+
+def _migrate_legacy_model_config(model_config: dict) -> dict:
+    """Translate a v0.2 checkpoint's flat `model_config` (giant/pipeline.py's
+    old shape: `hidden_dim`/`n_blocks`/`emb_dim`/`dropout`/`conditioning`/
+    `router`/`mode`/... all at one level) into the nested
+    `{"pdg_vocab", "mat_vocab", "conditioning", "stage1_model",
+    "stage2_model"}` shape `build_models` expects.
+
+    Sets `stage2_model.n_sec.owner = "stage1"` so the n_sec_head weights a v0.2
+    checkpoint carries on its Stage-1 module keep loading there instead of the new
+    default location (`Stage2OneShot`) — the n_sec head was trained against Stage 1's
+    own `ConditionEncoder` output, so it has to stay attached to Stage 1's module, not
+    just be labeled as such.
+
+    Only the monolithic (non-routed) trunk shape is exercised by the step-2
+    migration test; a routed v0.2 checkpoint still builds correctly here
+    (the router config passes through), but its
+    state dict isn't covered by `migrate_legacy_state_dict` below.
+    """
+    m = model_config
+    conditioning_mode = m.get("conditioning", "embedding")
+    generator = m.get("mode", "flow")
+    hidden_dim = m.get("hidden_dim", 256)
+    n_blocks = m.get("n_blocks", 6)
+    emb_dim = m.get("emb_dim", EMB_DIM)
+    dropout = m.get("dropout", 0.1)
+    k_max = m.get("k_max", K_MAX)
+    noise_dim = m.get("noise_dim", 64)
+    router_cfg = dict(m.get("router") or {})
+    reject_legacy_router_expert_sizing(router_cfg, source="this checkpoint's model_config.router")
+    router_cfg.setdefault("enabled", False)
+
+    F = V02_FIXED_FACTS
+    cond_n_layers = F["conditioning.particle.n_layers"]  # same fact for both axes
+    return {
+        "pdg_vocab": m["pdg_vocab"],
+        "mat_vocab": m["mat_vocab"],
+        "conditioning": {
+            "out_dim": F["conditioning.out_dim"],
+            "share_stages": False,
+            "particle": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": cond_n_layers},
+            "material": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": cond_n_layers},
+        },
+        "stage1_model": {
+            "active": F["stage1_model.active"],
+            "generator": generator,
+            "hidden_dim": hidden_dim,
+            "n_res_blocks": n_blocks,
+            "dropout": dropout,
+            "flow": {"time_dim": F["stage1_model.flow.time_dim"]},
+            "ddpm": {"time_dim": F["stage1_model.ddpm.time_dim"]},
+            "wgan": {"noise_dim": noise_dim},
+            "router": dict(router_cfg),
+        },
+        "stage2_model": {
+            "active": F["stage2_model.active"],
+            "decoder": F["stage2_model.decoder"],
+            "generator": generator,
+            "hidden_dim": hidden_dim,
+            "n_res_blocks": n_blocks,
+            "dropout": dropout,
+            "k_max": k_max,
+            "context_dim": F["stage2_model.context_dim"],
+            "n_sec": {"mode": "head", "owner": "stage1"},
+            "particle_type": {"target": F["stage2_model.particle_type.target"]},
+            "flow": {"time_dim": F["stage2_model.flow.time_dim"]},
+            "ddpm": {"time_dim": F["stage2_model.ddpm.time_dim"]},
+            "wgan": {"noise_dim": noise_dim},
+            "router": {**router_cfg, "tie_to_stage1": False},
+        },
+    }
+
+
+def migrate_legacy_state_dict(old_stage1_sd: dict, old_stage2_sd: dict) -> tuple[dict, dict]:
+    """Remap a v0.2 checkpoint's (`DenoisingMLP`-or-`WGANGenerator`,
+    `SecondaryDecoder`-or-`WGANSecondaryGenerator`) state dicts onto the new
+    `(Stage1Model, Stage2OneShot)` module structure produced by
+    `build_models(_migrate_legacy_model_config(model_config))`.
+
+    Only the monolithic (non-routed) trunk shape is handled.
+    """
+
+    def _trunk_prefix(k: str) -> str:
+        if k.startswith(("input_proj.", "blocks.", "out_proj.")):
+            return f"trunk.{k}"
+        return k
+
+    new_stage1 = {}
+    for k, v in old_stage1_sd.items():
+        if k.startswith("n_sec_head."):
+            new_stage1[k] = v  # stays top-level (n_sec.owner="stage1")
+        else:
+            new_stage1[_trunk_prefix(k)] = v
+
+    new_stage2 = {}
+    for k, v in old_stage2_sd.items():
+        if k.startswith("cond_enc.base."):
+            new_stage2["cond_enc." + k[len("cond_enc.base.") :]] = v
+        elif k.startswith("cond_enc.stage1_proj."):
+            new_stage2["context_adapter.proj." + k[len("cond_enc.stage1_proj.") :]] = v
+        elif k.startswith("cond_enc.fuse."):
+            new_stage2["fuse." + k[len("cond_enc.fuse.") :]] = v
+        else:
+            new_stage2[_trunk_prefix(k)] = v
+
+    return new_stage1, new_stage2
@@ -0,0 +1,234 @@
+"""Factories: `build_models`/`build_critics` assemble the top-level stage
+models from a config dict (issues.md Issue 8)."""
+
+import torch.nn as nn
+
+from giant.config import ConditioningConfig, Stage1ModelConfig, Stage2ModelConfig
+from giant.constants import X_DIM
+from giant.model._legacy import _migrate_legacy_model_config
+from giant.model.encoders import ConditionEncoder
+from giant.model.models import (
+    CriticModel,
+    Stage1Model,
+    Stage2Autoregressive,
+    Stage2OneShot,
+    resolve_type_n_classes,
+    stage2_trunk_sec_dim,
+)
+from giant.model.objectives import build_objective
+from giant.model.routers import Router, _build_router_from_cfg
+
+# ---------------------------------------------------------------------------
+# Factories
+# ---------------------------------------------------------------------------
+
+
+def build_models(model_config: dict) -> dict[str, nn.Module | None]:
+    """Construct `{"stage1": ..., "stage2": ...}` from a config dict — either
+    the new nested shape (has a `"stage1_model"` key, plus `"pdg_vocab"`/
+    `"mat_vocab"`/`"conditioning"` at the top level) or a v0.2 checkpoint's
+    flat `model_config`, auto-migrated via `_migrate_legacy_model_config`.
+
+    A stage is `None` in the result when that stage's `active = False`.
+    `stage2_model.router.tie_to_stage1` shares stage 1's literal `Router`
+    instance rather than building a second, independently-parameterized one
+    (v0.2's actual — probably accidental — behaviour: two routers built from
+    one config with no semantic relationship between them).
+
+    `conditioning.share_stages = true` builds one `ConditionEncoder`
+    instance here and passes it to both stages (`Stage1Model`/`Stage2OneShot`/
+    `Stage2Autoregressive`'s `cond_enc` param), instead of each stage
+    building its own — halving the conditioning parameter count and forcing a
+    common representation. `false` (default) keeps v0.2 behaviour:
+    independent instances with identical config but independent weights.
+    """
+    cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config)
+    pdg_vocab = cfg["pdg_vocab"]
+    mat_vocab = cfg["mat_vocab"]
+    conditioning_cfg = ConditioningConfig.from_dict(cfg["conditioning"])
+    particle_cfg = conditioning_cfg.particle
+    material_cfg = conditioning_cfg.material
+    particle_conditioning = particle_cfg.type
+    s1_spec = Stage1ModelConfig.from_dict(cfg["stage1_model"])
+    s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"])
+    cond_out_dim = conditioning_cfg.out_dim
+    shared_cond_enc: ConditionEncoder | None = None
+    if conditioning_cfg.share_stages:
+        shared_cond_enc = ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
+
+    result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None}
+
+    stage1_router: Router | None = None
+    if s1_spec.active:
+        router_cfg = cfg["stage1_model"].get("router") or {}
+        if s1_spec.router.enabled:
+            stage1_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning)
+        generator = s1_spec.generator
+        objective = build_objective(generator)
+        # wgan has no time_dim concept (no diffusion/flow time variable) —
+        # matches the pre-dataclass .get("time_dim", 64) fallback, which
+        # always hit its default for a wgan sub-block too.
+        time_dim = getattr(s1_spec, generator).time_dim if objective.needs_time else 64
+        n_sec_owner = s2_spec.n_sec.owner
+        n_sec_head_k_max = s2_spec.k_max if n_sec_owner == "stage1" else None
+        result["stage1"] = Stage1Model(
+            pdg_vocab=pdg_vocab,
+            mat_vocab=mat_vocab,
+            particle_cfg=particle_cfg,
+            material_cfg=material_cfg,
+            hidden_dim=s1_spec.hidden_dim,
+            n_res_blocks=s1_spec.n_res_blocks,
+            cond_out_dim=cond_out_dim,
+            dropout=s1_spec.dropout,
+            generator=generator,
+            time_dim=time_dim,
+            noise_dim=s1_spec.wgan.noise_dim,
+            router=stage1_router,
+            trunk_type=s1_spec.trunk.type,
+            block_conditioning=s1_spec.trunk.block_conditioning,
+            n_sec_head_k_max=n_sec_head_k_max,
+            cond_enc=shared_cond_enc,
+            n_sec_head_cfg=s1_spec.heads.n_sec.to_dict(),
+        )
+
+    if s2_spec.active:
+        decoder = s2_spec.decoder
+        router_cfg = cfg["stage2_model"].get("router") or {}
+        stage2_router: Router | None = None
+        if s2_spec.router.enabled:
+            if s2_spec.router.tie_to_stage1 and stage1_router is not None:
+                stage2_router = stage1_router
+            else:
+                stage2_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning)
+        generator = s2_spec.generator
+        objective = build_objective(generator)
+        # wgan has no time_dim concept — see the matching comment in stage 1
+        # above.
+        time_dim = getattr(s2_spec, generator).time_dim if objective.needs_time else 64
+        n_sec_owner = s2_spec.n_sec.owner
+        stop_token = s2_spec.n_sec.mode == "stop_token"
+        k_max = s2_spec.k_max
+        particle_type_cfg = s2_spec.particle_type
+
+        if decoder == "autoregressive":
+            ar_cfg = s2_spec.autoregressive
+            result["stage2"] = Stage2Autoregressive(
+                pdg_vocab=pdg_vocab,
+                mat_vocab=mat_vocab,
+                particle_cfg=particle_cfg,
+                material_cfg=material_cfg,
+                hidden_dim=s2_spec.hidden_dim,
+                n_res_blocks=s2_spec.n_res_blocks,
+                cond_out_dim=cond_out_dim,
+                context_dim=s2_spec.context_dim,
+                dropout=s2_spec.dropout,
+                generator=generator,
+                time_dim=time_dim,
+                noise_dim=s2_spec.wgan.noise_dim,
+                k_max=k_max,
+                router=stage2_router,
+                trunk_type=s2_spec.trunk.type,
+                block_conditioning=s2_spec.trunk.block_conditioning,
+                build_n_sec_head=n_sec_owner != "stage1" and not stop_token,
+                particle_type_cfg=particle_type_cfg,
+                history=ar_cfg.history,
+                attn_n_heads=ar_cfg.attn_n_heads,
+                attn_n_layers=ar_cfg.attn_n_layers,
+                cond_enc=shared_cond_enc,
+                n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
+                type_head_cfg=s2_spec.heads.type.to_dict(),
+                build_stop_head=stop_token,
+                stop_sampling=s2_spec.n_sec.stop_sampling,
+                stop_head_cfg=s2_spec.heads.n_sec.to_dict(),
+            )
+        else:
+            sec_dim = stage2_trunk_sec_dim(
+                particle_type_cfg, generator, k_max, resolve_type_n_classes(particle_type_cfg, particle_cfg.emb_dim)
+            )
+            result["stage2"] = Stage2OneShot(
+                pdg_vocab=pdg_vocab,
+                mat_vocab=mat_vocab,
+                particle_cfg=particle_cfg,
+                material_cfg=material_cfg,
+                hidden_dim=s2_spec.hidden_dim,
+                n_res_blocks=s2_spec.n_res_blocks,
+                cond_out_dim=cond_out_dim,
+                context_dim=s2_spec.context_dim,
+                sec_dim=sec_dim,
+                dropout=s2_spec.dropout,
+                generator=generator,
+                time_dim=time_dim,
+                noise_dim=s2_spec.wgan.noise_dim,
+                k_max=k_max,
+                router=stage2_router,
+                trunk_type=s2_spec.trunk.type,
+                block_conditioning=s2_spec.trunk.block_conditioning,
+                build_n_sec_head=n_sec_owner != "stage1",
+                particle_type_cfg=particle_type_cfg,
+                cond_enc=shared_cond_enc,
+                n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
+                type_head_cfg=s2_spec.heads.type.to_dict(),
+            )
+
+    return result
+
+
+def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
+    """Construct `{"stage1": ..., "stage2": ...}` critics for `generator =
+    "wgan"` training. Training-only — never persisted for inference the way
+    `build_models`'s pair is. `None` for a stage that's inactive or not
+    WGAN."""
+    cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config)
+    pdg_vocab = cfg["pdg_vocab"]
+    mat_vocab = cfg["mat_vocab"]
+    conditioning_cfg = ConditioningConfig.from_dict(cfg["conditioning"])
+    particle_cfg = conditioning_cfg.particle
+    material_cfg = conditioning_cfg.material
+    cond_out_dim = conditioning_cfg.out_dim
+    s1_spec = Stage1ModelConfig.from_dict(cfg["stage1_model"])
+    s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"])
+
+    result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None}
+
+    if s1_spec.active and build_objective(s1_spec.generator).is_adversarial:
+        result["stage1"] = CriticModel(
+            pdg_vocab=pdg_vocab,
+            mat_vocab=mat_vocab,
+            particle_cfg=particle_cfg,
+            material_cfg=material_cfg,
+            in_dim=X_DIM,
+            hidden_dim=s1_spec.wgan.critic_hidden_dim or s1_spec.hidden_dim,
+            n_res_blocks=s1_spec.wgan.critic_n_res_blocks or s1_spec.n_res_blocks,
+            cond_out_dim=cond_out_dim,
+            dropout=s1_spec.dropout,
+            stage="stage1",
+            trunk_type=s1_spec.trunk.type,
+            block_conditioning=s1_spec.trunk.block_conditioning,
+        )
+
+    if s2_spec.active and build_objective(s2_spec.generator).is_adversarial:
+        k_max = s2_spec.k_max
+        particle_type_cfg = s2_spec.particle_type
+        in_dim = stage2_trunk_sec_dim(
+            particle_type_cfg,
+            s2_spec.generator,
+            k_max,
+            resolve_type_n_classes(particle_type_cfg, particle_cfg.emb_dim),
+        )
+        result["stage2"] = CriticModel(
+            pdg_vocab=pdg_vocab,
+            mat_vocab=mat_vocab,
+            particle_cfg=particle_cfg,
+            material_cfg=material_cfg,
+            in_dim=in_dim,
+            hidden_dim=s2_spec.wgan.critic_hidden_dim or s2_spec.hidden_dim,
+            n_res_blocks=s2_spec.wgan.critic_n_res_blocks or s2_spec.n_res_blocks,
+            cond_out_dim=cond_out_dim,
+            dropout=s2_spec.dropout,
+            stage="stage2",
+            context_dim=s2_spec.context_dim,
+            trunk_type=s2_spec.trunk.type,
+            block_conditioning=s2_spec.trunk.block_conditioning,
+        )
+
+    return result
@@ -0,0 +1,101 @@
+"""Conditioning encoder — fuses continuous conditioning with particle/material
+identity (issues.md Issue 8)."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from giant.cond_layout import CondLayout
+from giant.config import ConditioningAxisConfig
+from giant.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM
+from giant.model.layers import _make_axis_mlp
+
+
+class ConditionEncoder(nn.Module):
+    """Fuses continuous conditioning with particle/material identity.
+
+    The particle and material axes are configured independently
+    (`particle_cfg`/`material_cfg`, each a `ConditioningAxisConfig`) and may
+    mix freely, e.g. material "physical" with particle "embedding". Three
+    modes per axis:
+      - "embedding": a learned `nn.Embedding` lookup, indexed by `cond_cat`'s
+        dense training-vocab index. Memorizes the training menu.
+      - "physical": an `n_layers`-deep MLP over the axis's raw physical
+        properties (already present in `cond_cont`'s physical block — see
+        giant.data.transforms.build_features), computable for any PDG code /
+        material name rather than only ones seen in training.
+      - "onehot": a fixed, unlearned one-hot vector over a top-N-plus-other
+        class map (`giant.data.loader.build_topn_map_from_files`/
+        `build_pdg_topn_map_from_files`), read from `cond_cat`'s extra
+        top-N-index column(s).
+
+    Every column index/slice comes from `self.layout`
+    (`giant.cond_layout.CondLayout`), the same object the feature builders
+    lay the arrays out with, so the two sides cannot drift apart.
+    """
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        particle_cfg: ConditioningAxisConfig,
+        material_cfg: ConditioningAxisConfig,
+        cont_dim: int = COND_DIM,
+        out_dim: int = 128,
+    ) -> None:
+        super().__init__()
+        self.particle_cfg = particle_cfg
+        self.material_cfg = material_cfg
+        # Also validates both axis types — an unknown one raises here.
+        self.layout = CondLayout.from_types(particle_cfg.type, material_cfg.type)
+
+        p_type = particle_cfg.type
+        p_emb_dim = particle_cfg.emb_dim
+        if p_type == "embedding":
+            self.pdg_emb = nn.Embedding(pdg_vocab, p_emb_dim)
+        elif p_type == "physical":
+            self.particle_mlp = _make_axis_mlp(PARTICLE_PHYS_DIM, p_emb_dim, particle_cfg.n_layers)
+
+        m_type = material_cfg.type
+        m_emb_dim = material_cfg.emb_dim
+        if m_type == "embedding":
+            self.mat_emb = nn.Embedding(mat_vocab, m_emb_dim)
+        elif m_type == "physical":
+            self.material_mlp = _make_axis_mlp(MATERIAL_PHYS_DIM, m_emb_dim, material_cfg.n_layers)
+
+        in_dim = COND_DIM_BASE + p_emb_dim + m_emb_dim
+        self.mlp = nn.Sequential(
+            nn.Linear(in_dim, out_dim),
+            nn.SiLU(),
+            nn.Linear(out_dim, out_dim),
+        )
+
+    def _particle_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor):
+        p_type = self.particle_cfg.type
+        if p_type == "embedding":
+            return self.pdg_emb(cond_cat[:, self.layout.PDG_COL])
+        if p_type == "physical":
+            return self.particle_mlp(cond_cont[:, self.layout.particle_phys])
+        assert self.layout.particle_topn_col is not None
+        return F.one_hot(
+            cond_cat[:, self.layout.particle_topn_col],
+            num_classes=self.particle_cfg.emb_dim,
+        ).float()
+
+    def _material_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor):
+        m_type = self.material_cfg.type
+        if m_type == "embedding":
+            return self.mat_emb(cond_cat[:, self.layout.MAT_COL])
+        if m_type == "physical":
+            return self.material_mlp(cond_cont[:, self.layout.material_phys])
+        assert self.layout.material_topn_col is not None
+        return F.one_hot(
+            cond_cat[:, self.layout.material_topn_col],
+            num_classes=self.material_cfg.emb_dim,
+        ).float()
+
+    def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        pdg_e = self._particle_embed(cond_cont, cond_cat)
+        mat_e = self._material_embed(cond_cont, cond_cat)
+        x = torch.cat([cond_cont[:, self.layout.base], pdg_e, mat_e], dim=-1)
+        return self.mlp(x)
@@ -0,0 +1,217 @@
+"""History encoders — stage-2 autoregressive only. Self-contained, no
+dependency on any other `giant.model` submodule (issues.md Issue 8), except
+for the `HISTORY_REGISTRY`/`build_history` factory, which mirrors
+`giant.model.routers`'s `Router`/`ROUTER_REGISTRY` pattern (gitea #35)."""
+
+import inspect
+
+import torch
+import torch.nn as nn
+
+
+class HistoryEncoder(nn.Module):
+    """Interface for stage-2 autoregressive per-token history summaries:
+    `forward(feat, has_prev) -> (B, K, out_dim)`, a single parallel pass over
+    a full (teacher-forced) token sequence — used by training. `MarkovHistory`
+    and `AttentionHistory` are the two registered implementations (see
+    `HISTORY_REGISTRY`/`build_history`). Inference (`giant/sample.py`)
+    generates one token at a time and cannot afford `forward`'s per-step cost
+    to be O(K) (attention would then be O(K^2) over a rollout's k_max loop),
+    so this interface also declares `init_cache`/`step` for that incremental
+    path, with working O(1) defaults here (`init_cache` -> `None`, `step` ->
+    one `forward` call ignoring `cache`) — correct for any encoder whose
+    per-step cost is already O(1) (i.e. it only ever looks at the previous
+    token, not the full prefix), which is what `MarkovHistory` relies on.
+    `AttentionHistory` overrides both with real incremental-cache versions,
+    since its `forward` genuinely needs the full prefix."""
+
+    def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
+
+    def init_cache(self) -> object:
+        return None
+
+    def step(self, feat: torch.Tensor, has_prev: torch.Tensor, cache: object) -> tuple[torch.Tensor, object]:
+        return self.forward(feat, has_prev), cache
+
+
+HISTORY_REGISTRY: dict[str, type[HistoryEncoder]] = {}
+
+
+def register_history(name: str):
+    def decorator(cls: type[HistoryEncoder]) -> type[HistoryEncoder]:
+        HISTORY_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def build_history(name: str, in_dim: int, out_dim: int, **kwargs) -> HistoryEncoder:
+    """Factory: look up a `HistoryEncoder` subclass by name from the registry.
+
+    Every registered history type is fed the same `stage2_model.autoregressive`
+    kwargs; kwargs not declared by that type's constructor are silently
+    dropped, so per-type hyperparameters (e.g. `AttentionHistory`'s
+    `n_heads`/`n_layers`) can coexist in one config without special-casing —
+    mirrors `giant.model.routers.build_router`.
+    """
+    if name not in HISTORY_REGISTRY:
+        raise ValueError(f"unknown history type {name!r}; available: {sorted(HISTORY_REGISTRY)}")
+    cls = HISTORY_REGISTRY[name]
+    accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "in_dim", "out_dim"}
+    filtered = {k: v for k, v in kwargs.items() if k in accepted}
+    return cls(in_dim, out_dim, **filtered)
+
+
+@register_history("none")
+class NoHistory(HistoryEncoder):
+    """No history signal at all — ignores feat/has_prev entirely and always
+    returns zeros. Ablates whether the AR decoder's history conditioning is
+    earning its parameters. `init_cache`/`step` use the base class's O(1)
+    defaults unmodified (this encoder's own `forward` is already O(1) per
+    call regardless of prefix length)."""
+
+    def __init__(self, in_dim: int, out_dim: int) -> None:
+        super().__init__()
+        self.out_dim = out_dim
+
+    def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
+        B, K, _ = feat.shape
+        return torch.zeros(B, K, self.out_dim, device=feat.device, dtype=feat.dtype)
+
+
+@register_history("markov")
+class MarkovHistory(HistoryEncoder):
+    """Summarizes the previous secondary's own `(energy_fraction, direction,
+    type_representation)` through one small MLP — the "markov" history:
+    token i+1 only ever sees token i plus the running scalars
+    (`remaining_frac`/`slot_idx`, fused in separately by
+    `Stage2Autoregressive._token_cond`), not the full prefix.
+
+    At slot 0 (`has_prev` False) substitutes a learned start vector rather
+    than zeros — a reasonable default.
+    """
+
+    def __init__(self, in_dim: int, out_dim: int) -> None:
+        super().__init__()
+        self.start = nn.Parameter(torch.zeros(in_dim))
+        self.mlp = nn.Sequential(nn.Linear(in_dim, out_dim), nn.SiLU())
+
+    def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
+        start = self.start.view(1, 1, -1).expand_as(feat)
+        x = torch.where(has_prev.unsqueeze(-1), feat, start)
+        return self.mlp(x)
+
+
+class _CausalAttnBlock(nn.Module):
+    """One pre-norm causal self-attention block for `AttentionHistory`.
+
+    Exposes two forward paths that must agree (see
+    `test_attention_history_step_matches_forward` in `tests/test_network.py`):
+    `forward` — the full-sequence, causally-masked pass used for training;
+    `step` — an incremental pass for inference, given the *pre-attention*
+    normalized hidden states of every earlier position (`kv_cache`, i.e.
+    `norm1(x)` for positions `< t`, not `x` itself). Caching `norm1(x)` rather
+    than raw `x` is what makes `step` correct: this block's attention needs
+    exactly that quantity as keys/values, and `LayerNorm` has no cross-position
+    interaction, so recomputing it per position instead of caching it would
+    still be correct but pointlessly repeat work. The *next* block's cache is
+    built from a different sequence (this block's output), so each block owns
+    an independent cache entry.
+    """
+
+    def __init__(self, dim: int, n_heads: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn = nn.MultiheadAttention(dim, n_heads, dropout=dropout, batch_first=True)
+        self.norm2 = nn.LayerNorm(dim)
+        self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))
+
+    def forward(self, x: torch.Tensor, causal_mask: torch.Tensor) -> torch.Tensor:
+        h = self.norm1(x)
+        attn_out, _ = self.attn(h, h, h, attn_mask=causal_mask, need_weights=False)
+        x = x + attn_out
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+    def step(self, x_new: torch.Tensor, kv_cache: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor]:
+        """`x_new`: `(B, 1, dim)`, this position's input. `kv_cache`: `None`
+        (first position) or `(B, T, dim)` — `norm1(x)` of every earlier
+        position at this same block. Returns `(out, new_kv_cache)`, `out`
+        being this position's block output (`(B, 1, dim)`, to feed the next
+        block's `step`), `new_kv_cache` the same cache extended by this
+        position (to reuse at this block's *next* `step` call)."""
+        h_new = self.norm1(x_new)
+        kv = h_new if kv_cache is None else torch.cat([kv_cache, h_new], dim=1)
+        attn_out, _ = self.attn(h_new, kv, kv, need_weights=False)
+        x = x_new + attn_out
+        x = x + self.mlp(self.norm2(x))
+        return x, kv
+
+
+@register_history("attention")
+class AttentionHistory(HistoryEncoder):
+    """Causal self-attention over the emitted-token prefix — the more
+    expressive alternative to `MarkovHistory`'s fixed previous-token-only
+    summary. `feat`/`has_prev`
+    follow the same shifted-by-one convention `MarkovHistory` and
+    `Stage2Autoregressive._token_cond` use: `feat[:, i]` is token `i - 1`'s
+    own `(energy_fraction, direction, type_representation)`, with a learned
+    start vector substituted at `has_prev == False` positions (only slot 0 in
+    practice — see `giant.training.stage2_inputs._ar_has_prev`). Causal masking then makes
+    position `i`'s output a function of `feat[:, 1:i+1]` — i.e. tokens
+    `0..i-1` — exactly the prefix available when predicting token `i`.
+
+    `forward` is the parallel training path (one pass over the whole
+    teacher-forced sequence); `init_cache`/`step` are the incremental
+    inference path `giant/sample.py` uses, one new token per call, to avoid
+    re-encoding the whole prefix from scratch every slot — `step` must be
+    called exactly once per slot (its cache-extension is not idempotent),
+    so a slot's output must be reused for
+    every model call within that slot (`forward`'s ODE substeps, or a separate
+    `predict_type` call) rather than re-derived — see
+    `Stage2Autoregressive.history_step`.
+    """
+
+    def __init__(self, in_dim: int, out_dim: int, n_heads: int = 4, n_layers: int = 2) -> None:
+        super().__init__()
+        self.start = nn.Parameter(torch.zeros(in_dim))
+        self.in_proj = nn.Linear(in_dim, out_dim)
+        self.blocks = nn.ModuleList([_CausalAttnBlock(out_dim, n_heads) for _ in range(n_layers)])
+
+    def _embed(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
+        start = self.start.view(1, 1, -1).expand_as(feat)
+        x = torch.where(has_prev.unsqueeze(-1), feat, start)
+        return self.in_proj(x)
+
+    def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
+        B, K, _ = feat.shape
+        x = self._embed(feat, has_prev)
+        mask = nn.Transformer.generate_square_subsequent_mask(K, device=feat.device)
+        for block in self.blocks:
+            x = block(x, mask)
+        return x
+
+    def init_cache(self) -> list[torch.Tensor | None]:
+        return [None for _ in self.blocks]
+
+    def step(
+        self,
+        feat: torch.Tensor,
+        has_prev: torch.Tensor,
+        cache: object,
+    ) -> tuple[torch.Tensor, object]:
+        """`feat`/`has_prev`: `(B, 1, in_dim)`/`(B, 1)` — the newest token's
+        own features (what would be `feat[:, k]` in `forward`). `cache`: the
+        `list[Tensor | None]` from `init_cache`/a previous `step` call (typed
+        `object` here to match `HistoryEncoder.step`'s base signature).
+        Advances every block's cache by this position and returns this
+        position's output (`(B, 1, out_dim)`, the correct history summary for
+        the NEXT slot) plus the updated cache."""
+        assert isinstance(cache, list)
+        x = self._embed(feat, has_prev)
+        new_cache: list[torch.Tensor | None] = []
+        for block, kv in zip(self.blocks, cache):
+            x, kv_new = block.step(x, kv)
+            new_cache.append(kv_new)
+        return x, new_cache
@@ -0,0 +1,184 @@
+"""Small stateless-ish building blocks shared across encoders/trunks/models —
+no dependency on any other `giant.model` submodule (issues.md Issue 8)."""
+
+import math
+
+import torch
+import torch.nn as nn
+
+
+class SinusoidalEmbedding(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        assert dim % 2 == 0, "dim must be even"
+        half = dim // 2
+        freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32) / max(half - 1, 1))
+        self.register_buffer("freqs", freqs)
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        t = t.reshape(-1, 1).float()
+        args = t * self.freqs.unsqueeze(0)  # (B, half)
+        return torch.cat([args.sin(), args.cos()], dim=-1)  # (B, dim)
+
+
+def _make_axis_mlp(in_dim: int, emb_dim: int, n_layers: int) -> nn.Sequential:
+    """`n_layers`-deep MLP producing an `emb_dim`-wide vector from `in_dim`
+    physical properties (`conditioning.{particle,material}.n_layers`).
+
+    `n_layers=1` (the v0.3.0 default): a single `Linear`, no hidden
+    activation. `n_layers=2` reproduces v0.2's hardcoded depth exactly —
+    `Linear -> SiLU -> Linear` — which is why `migrate_config` back-fills
+    `n_layers=2` for migrated configs rather than the v0.3 default of 1 (see
+    its docstring).
+    """
+    if n_layers < 1:
+        raise ValueError(f"n_layers must be >= 1, got {n_layers}")
+    if n_layers == 1:
+        return nn.Sequential(nn.Linear(in_dim, emb_dim))
+    layers: list[nn.Module] = [nn.Linear(in_dim, emb_dim), nn.SiLU()]
+    for _ in range(n_layers - 2):
+        layers += [nn.Linear(emb_dim, emb_dim), nn.SiLU()]
+    layers.append(nn.Linear(emb_dim, emb_dim))
+    return nn.Sequential(*layers)
+
+
+def build_mlp_head(
+    in_dim: int, out_dim: int, hidden: int, depth: int = 2, act: type[nn.Module] = nn.SiLU
+) -> nn.Sequential:
+    """`depth`-layer MLP head (gitea #36) — factors out the n_sec_head/
+    type_head pattern duplicated five times across `giant.model.models`.
+
+    `depth=1` is a bare `Linear(in_dim, out_dim)` (no hidden layer/
+    activation); `depth>=2` is `Linear(in_dim, hidden) -> act -> [Linear
+    (hidden, hidden) -> act] * (depth-2) -> Linear(hidden, out_dim)` —
+    `depth=2` reproduces every pre-#36 n_sec_head/type_head exactly when
+    `hidden == hidden_dim // 2`. Mirrors `_make_axis_mlp`'s depth
+    convention above, but takes `hidden` and `out_dim` as independent
+    widths (n_sec_head/type_head's hidden width is not their output width,
+    unlike the particle/material axis MLPs)."""
+    if depth < 1:
+        raise ValueError(f"depth must be >= 1, got {depth}")
+    if depth == 1:
+        return nn.Sequential(nn.Linear(in_dim, out_dim))
+    layers: list[nn.Module] = [nn.Linear(in_dim, hidden), act()]
+    for _ in range(depth - 2):
+        layers += [nn.Linear(hidden, hidden), act()]
+    layers.append(nn.Linear(hidden, out_dim))
+    return nn.Sequential(*layers)
+
+
+class ContextAdapter(nn.Module):
+    """Projects a stage's outcome (e.g. Stage 1's 9D target) down to a
+    fixed-width context vector for a downstream stage's conditioning —
+    `stage2_model.context_dim`. Was `SecondaryConditionEncoder.stage1_proj`
+    (+ its `tanh`) in v0.2; pulled out as its own module in v0.3.0 since
+    `SecondaryConditionEncoder` as a wrapper class disappears."""
+
+    def __init__(self, in_dim: int, context_dim: int) -> None:
+        super().__init__()
+        self.proj = nn.Linear(in_dim, context_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.tanh(self.proj(x))
+
+
+BLOCK_REGISTRY: dict[str, type[nn.Module]] = {}
+
+
+def register_block(name: str):
+    def decorator(cls: type[nn.Module]) -> type[nn.Module]:
+        BLOCK_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def build_block(name: str, dim: int, cond_dim: int, dropout: float = 0.0) -> nn.Module:
+    """Factory: look up a registered conditioning-injection block by name and
+    construct one instance — `trunk.block_conditioning` (gitea #34)."""
+    if name not in BLOCK_REGISTRY:
+        raise ValueError(f"unknown block conditioning type {name!r}; available: {sorted(BLOCK_REGISTRY)}")
+    return BLOCK_REGISTRY[name](dim, cond_dim, dropout)
+
+
+@register_block("add")
+class ResBlock(nn.Module):
+    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.linear1 = nn.Linear(dim, dim)
+        self.cond_proj = nn.Linear(cond_dim, dim, bias=False)
+        self.act = nn.SiLU()
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        h = self.norm(x)
+        h = self.linear1(h) + self.cond_proj(cond)
+        h = self.act(h)
+        h = self.dropout(h)
+        h = self.linear2(h)
+        return x + h
+
+
+@register_block("film")
+class FilmResBlock(nn.Module):
+    """FiLM conditioning (Perez et al. 2018): a per-channel scale+shift
+    modulates the normalized features, on top of the norm's own affine —
+    an *additional* modulation, unlike `AdaLNResBlock` below, which replaces
+    the norm's affine outright. `film_proj` is zero-initialized so
+    `gamma=beta=0` at construction — conditioning has no effect on the
+    output until training moves it, a stable starting point (though not a
+    literal identity block, since `linear1`/`linear2` aren't zero-init)."""
+
+    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.linear1 = nn.Linear(dim, dim)
+        self.film_proj = nn.Linear(cond_dim, 2 * dim)
+        nn.init.zeros_(self.film_proj.weight)
+        nn.init.zeros_(self.film_proj.bias)
+        self.act = nn.SiLU()
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        h = self.norm(x)
+        gamma, beta = self.film_proj(cond).chunk(2, dim=-1)
+        h = h * (1 + gamma) + beta
+        h = self.linear1(h)
+        h = self.act(h)
+        h = self.dropout(h)
+        h = self.linear2(h)
+        return x + h
+
+
+@register_block("adaln")
+class AdaLNResBlock(nn.Module):
+    """AdaLN-Zero conditioning (DiT, Peebles & Xie 2022): the norm's own
+    affine is replaced by a conditioning-derived scale/shift, and the
+    residual branch is scaled by a conditioning-derived gate. `adaln_proj`
+    is zero-initialized, so `scale=shift=gate=0` at construction — the block
+    is the exact identity function at init (`x + 0 * h' == x`), regardless
+    of `x`/`cond`."""
+
+    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False)
+        self.linear1 = nn.Linear(dim, dim)
+        self.adaln_proj = nn.Linear(cond_dim, 3 * dim)
+        nn.init.zeros_(self.adaln_proj.weight)
+        nn.init.zeros_(self.adaln_proj.bias)
+        self.act = nn.SiLU()
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        h = self.norm(x)
+        scale, shift, gate = self.adaln_proj(cond).chunk(3, dim=-1)
+        h = h * (1 + scale) + shift
+        h = self.linear1(h)
+        h = self.act(h)
+        h = self.dropout(h)
+        h = self.linear2(h)
+        return x + gate * h
@@ -0,0 +1,783 @@
+"""Top-level stage models: `Stage1Model`, `Stage2OneShot`, `Stage2Autoregressive`,
+`CriticModel` — composed from encoders/trunks/history (issues.md Issue 8)."""
+
+import torch
+import torch.nn as nn
+
+from giant.config import ConditioningAxisConfig, HeadConfig, ParticleTypeConfig
+from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, SEC_SLOT_DIM, X_DIM
+from giant.model.encoders import ConditionEncoder
+from giant.model.history import HistoryEncoder, build_history
+from giant.model.layers import ContextAdapter, SinusoidalEmbedding, build_mlp_head
+from giant.model.objectives import build_objective
+from giant.model.routers import Router
+from giant.model.trunks import build_trunk
+
+# ---------------------------------------------------------------------------
+# Stage models
+# ---------------------------------------------------------------------------
+
+
+def resolve_type_n_classes(particle_type_cfg: ParticleTypeConfig, particle_emb_dim: int) -> int:
+    """Effective width fed to `stage2_type_dim`/`stage2_trunk_sec_dim` in
+    place of a bare `conditioning.particle.emb_dim` read. Under
+    `target = "onehot"` this is `stage2_model.particle_type.n_classes` (0 =
+    inherit `conditioning.particle.emb_dim`) — see gitea #29, which decoupled
+    the secondary-species vocabulary size from the unrelated
+    physical-conditioning MLP's output width. Under `target = "embedding"`
+    (or `"physical"`, which ignores this value entirely) `n_classes` doesn't
+    apply — the width stays `conditioning.particle.emb_dim`, the embedding
+    table's own dimensionality (`validate_config` requires
+    `conditioning.particle.type = "embedding"` here)."""
+    if particle_type_cfg.target == "onehot":
+        return particle_type_cfg.n_classes or particle_emb_dim
+    return particle_emb_dim
+
+
+def stage2_type_dim(particle_type_cfg: ParticleTypeConfig, emb_dim: int) -> int:
+    """Width of a single secondary slot's type slice —
+    `PARTICLE_PHYS_DIM` (log_mass, charge) for `target = "physical"`, else
+    `emb_dim` (both `"onehot"` class logits and `"embedding"` vectors are
+    this many classes/dims wide — callers resolve `emb_dim` via
+    `resolve_type_n_classes` first)."""
+    return PARTICLE_PHYS_DIM if particle_type_cfg.target == "physical" else emb_dim
+
+
+def stage2_trunk_sec_dim(particle_type_cfg: ParticleTypeConfig, generator: str, k_max: int, emb_dim: int) -> int:
+    """`Stage2OneShot`'s trunk output width.
+
+    `target = "physical"` is untouched from v0.2/today:
+    `k_max * SEC_SLOT_DIM`, the type slice folded into the same
+    flow-matched/WGAN vector as the continuous stick/dir slots.
+
+    `target` in `("onehot", "embedding")`: under an objective with
+    `folds_type_slice` (currently just wgan) the type slice is still folded
+    in (adversarial for onehot via ST-Gumbel, already-continuous for
+    embedding), just `emb_dim` wide instead of `PARTICLE_PHYS_DIM` wide:
+    `k_max * (CONT_SLOT_DIM + emb_dim)`. Otherwise (flow/ddpm) the type slice
+    isn't part of this vector at all — it's `Stage2OneShot.type_head`'s job
+    instead — so the trunk only covers `k_max * CONT_SLOT_DIM`.
+    """
+    if particle_type_cfg.target == "physical":
+        return k_max * SEC_SLOT_DIM
+    if build_objective(generator).folds_type_slice:
+        return k_max * (CONT_SLOT_DIM + emb_dim)
+    return k_max * CONT_SLOT_DIM
+
+
+class StageModel(nn.Module):
+    """Base owning the scaffolding common to `Stage1Model`, `Stage2OneShot`,
+    `Stage2Autoregressive` (gitea #39): build-or-share `cond_enc`,
+    `particle_type_cfg` normalisation, and — via `_build_trunk_and_heads`,
+    called by each subclass's `__init__` once its own conditioning-assembly
+    modules exist — the objective/time-embedding/trunk construction and the
+    `n_sec_head`/`type_head` classifier heads. A subclass supplies only its
+    own conditioning assembly (`Stage1Model` uses `cond_enc` directly;
+    `Stage2OneShot`/`Stage2Autoregressive` add a context-fusion path) and its
+    trunk's output width.
+
+    `cond_enc`, if given, is used in place of building a fresh
+    `ConditionEncoder` — `conditioning.share_stages = true`: `build_models`
+    constructs one shared instance and passes it to both stages, halving the
+    conditioning parameter count and forcing a common representation."""
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        particle_cfg: ConditioningAxisConfig,
+        material_cfg: ConditioningAxisConfig,
+        cond_out_dim: int,
+        generator: str,
+        noise_dim: int,
+        k_max: int | None = None,
+        particle_type_cfg: ParticleTypeConfig | None = None,
+        cond_enc: ConditionEncoder | None = None,
+    ) -> None:
+        super().__init__()
+        self.generator_kind = generator
+        self.noise_dim = noise_dim
+        self.k_max = k_max
+        # `ParticleTypeConfig()`'s own dataclass default is target="onehot"
+        # (the config.toml default when [stage2_model.particle_type] is
+        # omitted) — a different question from "nobody passed anything to
+        # this constructor", which direct/test construction relies on
+        # defaulting to "physical" (build_models/build_critics always pass
+        # particle_type_cfg explicitly, so this sentinel is never hit there).
+        self.particle_type_cfg = (
+            particle_type_cfg if particle_type_cfg is not None else ParticleTypeConfig(target="physical")
+        )
+        self.type_dim = stage2_type_dim(
+            self.particle_type_cfg, resolve_type_n_classes(self.particle_type_cfg, particle_cfg.emb_dim)
+        )
+        self.cond_enc = (
+            cond_enc
+            if cond_enc is not None
+            else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
+        )
+
+    def _build_trunk_and_heads(
+        self,
+        *,
+        trunk_out_dim: int,
+        hidden_dim: int,
+        n_res_blocks: int,
+        cond_out_dim: int,
+        time_dim: int,
+        router: Router | None,
+        trunk_type: str,
+        block_conditioning: str,
+        dropout: float,
+        n_sec_head_k_max: int | None,
+        n_sec_head_cfg: dict | None,
+        type_head_out_dim: int | None,
+        type_head_cfg: dict | None,
+        build_stop_head: bool = False,
+        stop_head_cfg: dict | None = None,
+    ) -> None:
+        """Builds `self.time_emb`, `self.trunk`, `self.n_sec_head`,
+        `self.type_head`, `self.stop_head`. Called by a subclass's `__init__`
+        after it has set up its own conditioning-assembly modules —
+        `merged_cond_dim` below must match the width that assembly
+        (`_cond_embed`/`_base_cond`/`_token_cond`, or plain `cond_enc` for
+        `Stage1Model`) actually produces.
+
+        `n_sec_head` is built iff `n_sec_head_k_max is not None` (output
+        width `n_sec_head_k_max + 1`) — `Stage1Model` passes this only for a
+        migrated v0.2 checkpoint, `Stage2OneShot`/`Stage2Autoregressive` pass
+        it whenever `build_n_sec_head=True`. `type_head` is built iff
+        `type_head_out_dim is not None` (the caller — only the two Stage2
+        classes — passes `None` exactly when `particle_type_cfg.target ==
+        "physical"`) *and* the objective doesn't fold the type slice into its
+        own trunk output (checked here, since `objective` is already needed
+        for the trunk itself). `stop_head` is built iff `build_stop_head` —
+        only `Stage2Autoregressive` ever passes `True` (`n_sec.mode ==
+        "stop_token"`, mutually exclusive with `n_sec_head`), a single
+        `cond_out_dim -> 1` logit per call, same `HeadConfig` shape rules as
+        the other two heads.
+        """
+        objective = build_objective(self.generator_kind)
+        has_time = objective.needs_time
+        self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
+        merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
+        in_dim = objective.trunk_in_dim(trunk_out_dim, self.noise_dim)
+        self.trunk = build_trunk(
+            router,
+            trunk_type,
+            in_dim,
+            trunk_out_dim,
+            hidden_dim,
+            n_res_blocks,
+            merged_cond_dim,
+            dropout,
+            block_conditioning,
+        )
+        self.n_sec_head = None
+        if n_sec_head_k_max is not None:
+            head_cfg = HeadConfig.from_dict(n_sec_head_cfg)
+            hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
+            self.n_sec_head = build_mlp_head(cond_out_dim, n_sec_head_k_max + 1, hidden, head_cfg.depth)
+        self.type_head = None
+        if type_head_out_dim is not None and not objective.folds_type_slice:
+            head_cfg = HeadConfig.from_dict(type_head_cfg)
+            hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
+            self.type_head = build_mlp_head(cond_out_dim, type_head_out_dim, hidden, head_cfg.depth)
+        self.stop_head = None
+        if build_stop_head:
+            head_cfg = HeadConfig.from_dict(stop_head_cfg)
+            hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
+            self.stop_head = build_mlp_head(cond_out_dim, 1, hidden, head_cfg.depth)
+
+    def _build_context_fusion(self, x_dim: int, context_dim: int, cond_out_dim: int) -> None:
+        """Builds `self.context_adapter`/`self.fuse` — the stage-2-style
+        context-fusion pattern (project the previous stage's outcome down to
+        `context_dim` via `ContextAdapter`, concat onto the base conditioning,
+        project back to `cond_out_dim`) shared by `Stage2OneShot` and a
+        `stage="stage2"` `CriticModel` (gitea #57). Call from a subclass's
+        `__init__` before using `_cond_embed`."""
+        self.context_adapter = ContextAdapter(x_dim, context_dim)
+        self.fuse = nn.Sequential(
+            nn.Linear(cond_out_dim + context_dim, cond_out_dim),
+            nn.SiLU(),
+        )
+
+    def _cond_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
+        """Fuses base conditioning with the previous stage's outcome — pairs
+        with `_build_context_fusion`."""
+        base = self.cond_enc(cond_cont, cond_cat)
+        ctx = self.context_adapter(stage1_out)
+        return self.fuse(torch.cat([base, ctx], dim=-1))
+
+    def _require_n_sec_head(self) -> None:
+        if self.n_sec_head is None:
+            raise RuntimeError(
+                f"this {type(self).__name__} has no n_sec_head — it belongs to "
+                "a migrated v0.2 checkpoint (n_sec.owner='stage1'); call "
+                "stage1.predict_n_sec(cond_cont, cond_cat) instead"
+            )
+
+    def _require_type_head(self) -> None:
+        if self.type_head is None:
+            raise RuntimeError(
+                f"this {type(self).__name__} has no type_head — either "
+                "particle_type.target='physical' (the type slice is part of "
+                "forward()'s own output) or generator='wgan' (the WGAN "
+                "trainer reads the type slice out of forward()'s output "
+                "directly instead)"
+            )
+
+    def _require_stop_head(self) -> None:
+        if self.stop_head is None:
+            raise RuntimeError(
+                f"this {type(self).__name__} has no stop_head — only a "
+                "Stage2Autoregressive built with stage2_model.n_sec.mode = "
+                "'stop_token' owns one"
+            )
+
+
+class Stage1Model(StageModel):
+    """Predicts the 9D primary post-step vector. No `n_sec_head` — fresh runs
+    move it to stage 2, except for a migrated v0.2 checkpoint
+    (`n_sec_head_k_max` given), where it stays attached here
+    since that's where its weights live and what conditioning it was trained
+    against (see `_migrate_legacy_model_config`)."""
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        particle_cfg: ConditioningAxisConfig,
+        material_cfg: ConditioningAxisConfig,
+        hidden_dim: int = 256,
+        n_res_blocks: int = 6,
+        cond_out_dim: int = 128,
+        x_dim: int = X_DIM,
+        dropout: float = 0.0,
+        generator: str = "flow",
+        time_dim: int = 64,
+        noise_dim: int = 64,
+        router: Router | None = None,
+        trunk_type: str = "resmlp",
+        block_conditioning: str = "add",
+        n_sec_head_k_max: int | None = None,
+        cond_enc: ConditionEncoder | None = None,
+        n_sec_head_cfg: dict | None = None,
+    ) -> None:
+        super().__init__(
+            pdg_vocab,
+            mat_vocab,
+            particle_cfg,
+            material_cfg,
+            cond_out_dim=cond_out_dim,
+            generator=generator,
+            noise_dim=noise_dim,
+            cond_enc=cond_enc,
+        )
+        self._build_trunk_and_heads(
+            trunk_out_dim=x_dim,
+            hidden_dim=hidden_dim,
+            n_res_blocks=n_res_blocks,
+            cond_out_dim=cond_out_dim,
+            time_dim=time_dim,
+            router=router,
+            trunk_type=trunk_type,
+            block_conditioning=block_conditioning,
+            dropout=dropout,
+            n_sec_head_k_max=n_sec_head_k_max,
+            n_sec_head_cfg=n_sec_head_cfg,
+            type_head_out_dim=None,
+            type_head_cfg=None,
+        )
+
+    def forward(
+        self,
+        x_t: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        t: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        c_emb = self.cond_enc(cond_cont, cond_cat)
+        cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb
+        return self.trunk(x_t, cond, cond_cont, cond_cat)
+
+    def _require_n_sec_head(self) -> None:
+        """Overrides `StageModel`'s guard — a `Stage1Model` with no
+        `n_sec_head` points the caller to stage 2 (n_sec's default owner),
+        not to `stage1` as the base's message would."""
+        if self.n_sec_head is None:
+            raise RuntimeError(
+                "this Stage1Model has no n_sec_head — n_sec now lives on "
+                "stage 2 by default; this method only exists "
+                "for a migrated v0.2 checkpoint (n_sec.owner='stage1')"
+            )
+
+    def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        """Return n_sec logits (B, K_MAX+1) from conditioning alone. Only
+        valid on a migrated v0.2 checkpoint's Stage1Model — fresh v0.3.0
+        configs predict n_sec from Stage2OneShot instead."""
+        self._require_n_sec_head()
+        assert self.n_sec_head is not None
+        c_emb = self.cond_enc(cond_cont, cond_cat)
+        return self.n_sec_head(c_emb)
+
+
+class Stage2OneShot(StageModel):
+    """Predicts all `k_max` secondary slots simultaneously — v0.2 behaviour,
+    reproduced exactly (`decoder = "autoregressive"` is `Stage2Autoregressive`,
+    step 4/5, not implemented yet).
+
+    Owns `n_sec_head` by default unless `build_n_sec_head=False`
+    (a migrated v0.2 checkpoint, whose n_sec_head instead attaches to
+    Stage1Model — see `_migrate_legacy_model_config`).
+
+    `particle_type_cfg.target` (default `"physical"`) selects the
+    secondary-type mechanism: `"physical"` keeps the type slice folded into
+    the trunk's own
+    flow-matched/WGAN output, unchanged from v0.2 (`sec_dim` — computed by
+    the caller via `stage2_trunk_sec_dim` — already reflects this). Under
+    `"onehot"`/`"embedding"` with an objective (`giant.model.objectives`) that
+    doesn't fold the type slice (flow/ddpm), the type
+    slice is predicted by a separate `type_head` instead (same shape pattern
+    as `n_sec_head`) — `sec_dim` then covers only the continuous
+    stick/dir slots, `type_head` covers `k_max * emb_dim` type logits/vectors.
+    Under a folding objective (wgan) the type slice stays folded into `sec_dim`
+    (just `emb_dim` instead of `PARTICLE_PHYS_DIM` wide) and `type_head` is
+    unused (`None`) — the WGAN trainer handles the ST-Gumbel relaxation.
+    """
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        particle_cfg: ConditioningAxisConfig,
+        material_cfg: ConditioningAxisConfig,
+        hidden_dim: int = 256,
+        n_res_blocks: int = 6,
+        cond_out_dim: int = 128,
+        context_dim: int = 64,
+        sec_dim: int = SEC_DIM,
+        x_dim: int = X_DIM,
+        dropout: float = 0.0,
+        generator: str = "wgan",
+        time_dim: int = 64,
+        noise_dim: int = 64,
+        k_max: int = K_MAX,
+        router: Router | None = None,
+        trunk_type: str = "resmlp",
+        block_conditioning: str = "add",
+        build_n_sec_head: bool = True,
+        particle_type_cfg: ParticleTypeConfig | None = None,
+        cond_enc: ConditionEncoder | None = None,
+        n_sec_head_cfg: dict | None = None,
+        type_head_cfg: dict | None = None,
+    ) -> None:
+        super().__init__(
+            pdg_vocab,
+            mat_vocab,
+            particle_cfg,
+            material_cfg,
+            cond_out_dim=cond_out_dim,
+            generator=generator,
+            noise_dim=noise_dim,
+            k_max=k_max,
+            particle_type_cfg=particle_type_cfg,
+            cond_enc=cond_enc,
+        )
+        self._build_context_fusion(x_dim, context_dim, cond_out_dim)
+        target = self.particle_type_cfg.target
+        type_head_out_dim = None if target == "physical" else k_max * self.type_dim
+        self._build_trunk_and_heads(
+            trunk_out_dim=sec_dim,
+            hidden_dim=hidden_dim,
+            n_res_blocks=n_res_blocks,
+            cond_out_dim=cond_out_dim,
+            time_dim=time_dim,
+            router=router,
+            trunk_type=trunk_type,
+            block_conditioning=block_conditioning,
+            dropout=dropout,
+            n_sec_head_k_max=k_max if build_n_sec_head else None,
+            n_sec_head_cfg=n_sec_head_cfg,
+            type_head_out_dim=type_head_out_dim,
+            type_head_cfg=type_head_cfg,
+        )
+
+    def forward(
+        self,
+        x_t: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+        t: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
+        cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb
+        return self.trunk(x_t, cond, cond_cont, cond_cat)
+
+    def predict_n_sec(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+    ) -> torch.Tensor:
+        self._require_n_sec_head()
+        assert self.n_sec_head is not None
+        c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
+        return self.n_sec_head(c_emb)
+
+    def predict_type(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+    ) -> torch.Tensor:
+        """`(B, k_max, emb_dim)` per-slot type logits (`target="onehot"`) or
+        vectors (`target="embedding"`) — only under `generator in ("flow",
+        "ddpm")`; `generator == "wgan"` folds the type slice into `forward`'s
+        own output instead (see class docstring)."""
+        self._require_type_head()
+        assert self.type_head is not None
+        c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
+        return self.type_head(c_emb).view(-1, self.k_max, self.type_dim)
+
+
+class Stage2Autoregressive(StageModel):
+    """Emits secondaries one at a time in descending-energy order, instead
+    of `Stage2OneShot`'s simultaneous
+    k_max-slot prediction. `history` selects `MarkovHistory` or
+    `AttentionHistory` (`attn_n_heads`/`attn_n_layers`, attention only).
+    `teacher_forcing` handling lives entirely in the trainer
+    (`giant/train.py`), since it only affects how training inputs are
+    assembled, not this module's architecture.
+
+    Under teacher forcing every token's conditioning is built from ground
+    truth, so a whole K-token sequence trains in one parallel batched pass:
+    `forward` accepts `(B, K, ...)` tensors for an arbitrary K (not hardcoded
+    to `k_max`) — this also means a future one-token-at-a-time inference loop
+    (`K=1` per call, step 6) needs no interface change here.
+
+    Two independent conditioning paths, mirroring `Stage2OneShot`'s
+    `_cond_embed` but split in two: `_base_cond` (`cond_enc` +
+    `context_adapter` only) feeds `predict_n_sec`, since n_sec doesn't depend
+    on token position; `_token_cond` additionally fuses in the history
+    encoding and two running scalars (remaining energy-budget fraction,
+    normalized slot index), and feeds `forward`/`predict_type`/`predict_stop`/
+    the trunk.
+
+    `n_sec.mode = "stop_token"` (`build_stop_head=True`) replaces
+    `predict_n_sec`'s one-shot classifier with `predict_stop`'s per-token EOS
+    logit instead — the two heads are mutually exclusive (`build_n_sec_head`
+    is `False` whenever this is `True`, see `giant.model.builders`).
+    """
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        particle_cfg: ConditioningAxisConfig,
+        material_cfg: ConditioningAxisConfig,
+        hidden_dim: int = 256,
+        n_res_blocks: int = 6,
+        cond_out_dim: int = 128,
+        context_dim: int = 64,
+        x_dim: int = X_DIM,
+        dropout: float = 0.0,
+        generator: str = "wgan",
+        time_dim: int = 64,
+        noise_dim: int = 64,
+        k_max: int = K_MAX,
+        router: Router | None = None,
+        trunk_type: str = "resmlp",
+        block_conditioning: str = "add",
+        build_n_sec_head: bool = True,
+        particle_type_cfg: ParticleTypeConfig | None = None,
+        history: str = "markov",
+        attn_n_heads: int = 4,
+        attn_n_layers: int = 2,
+        cond_enc: ConditionEncoder | None = None,
+        n_sec_head_cfg: dict | None = None,
+        type_head_cfg: dict | None = None,
+        build_stop_head: bool = False,
+        stop_sampling: str = "greedy",
+        stop_head_cfg: dict | None = None,
+    ) -> None:
+        super().__init__(
+            pdg_vocab,
+            mat_vocab,
+            particle_cfg,
+            material_cfg,
+            cond_out_dim=cond_out_dim,
+            generator=generator,
+            noise_dim=noise_dim,
+            k_max=k_max,
+            particle_type_cfg=particle_type_cfg,
+            cond_enc=cond_enc,
+        )
+        self.history_kind = history
+        self.stop_sampling = stop_sampling
+        self.context_adapter = ContextAdapter(x_dim, context_dim)
+        self.base_fuse = nn.Sequential(
+            nn.Linear(cond_out_dim + context_dim, cond_out_dim),
+            nn.SiLU(),
+        )
+
+        # Reuses conditioning.out_dim for the history encoder's own output
+        # width — there's no dedicated stage2_model.autoregressive key for
+        # this, a reasonable default rather than a design-doc-specified value.
+        history_dim = cond_out_dim
+        hist_in_dim = CONT_SLOT_DIM + self.type_dim
+        self.history_encoder: HistoryEncoder = build_history(
+            history, hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers
+        )
+        token_fuse_in = cond_out_dim + context_dim + history_dim + 2  # +2: remaining_frac, slot_idx
+        self.token_fuse = nn.Sequential(
+            nn.Linear(token_fuse_in, cond_out_dim),
+            nn.SiLU(),
+        )
+
+        # `self.type_dim` (set by StageModel.__init__) doubles as the raw
+        # `emb_dim` `stage2_trunk_sec_dim` wants: for a non-"physical" target
+        # `stage2_type_dim` already resolved `type_dim` to exactly that value;
+        # for "physical" the emb_dim argument goes unused anyway.
+        token_dim = stage2_trunk_sec_dim(self.particle_type_cfg, generator, 1, self.type_dim)
+        target = self.particle_type_cfg.target
+        type_head_out_dim = None if target == "physical" else self.type_dim
+        self._build_trunk_and_heads(
+            trunk_out_dim=token_dim,
+            hidden_dim=hidden_dim,
+            n_res_blocks=n_res_blocks,
+            cond_out_dim=cond_out_dim,
+            time_dim=time_dim,
+            router=router,
+            trunk_type=trunk_type,
+            block_conditioning=block_conditioning,
+            dropout=dropout,
+            n_sec_head_k_max=k_max if build_n_sec_head else None,
+            n_sec_head_cfg=n_sec_head_cfg,
+            type_head_out_dim=type_head_out_dim,
+            type_head_cfg=type_head_cfg,
+            build_stop_head=build_stop_head,
+            stop_head_cfg=stop_head_cfg,
+        )
+
+    def _base_cond(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
+        base = self.cond_enc(cond_cont, cond_cat)
+        ctx = self.context_adapter(stage1_out)
+        return self.base_fuse(torch.cat([base, ctx], dim=-1))
+
+    def _token_cond(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+        history_feat: torch.Tensor,
+        has_prev: torch.Tensor,
+        remaining_frac: torch.Tensor,
+        slot_idx: torch.Tensor,
+        hist: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """`hist`, if given, overrides recomputing `self.history_encoder`
+        from `history_feat`/`has_prev` — the inference-time KV-cache path
+        (`Stage2Autoregressive.history_step`) precomputes it once per slot and
+        passes it in here so a slot's (possibly several) model calls — an ODE
+        loop's substeps, or a separate `predict_type` call — read the same
+        cached history instead of each re-deriving (and, under attention,
+        re-appending to the cache — see `AttentionHistory.step`'s docstring)."""
+        K = history_feat.size(1)
+        base = self.cond_enc(cond_cont, cond_cat).unsqueeze(1).expand(-1, K, -1)
+        ctx = self.context_adapter(stage1_out).unsqueeze(1).expand(-1, K, -1)
+        if hist is None:
+            hist = self.history_encoder(history_feat, has_prev)
+        scalars = torch.stack([remaining_frac, slot_idx], dim=-1)
+        return self.token_fuse(torch.cat([base, ctx, hist, scalars], dim=-1))
+
+    def init_history_cache(self):
+        """Inference-only incremental-decoding state for `self.history_encoder`
+        (`giant/sample.py`'s AR loop) — whatever `self.history_encoder.init_cache()`
+        returns for the configured `history` type: `None` under `history="markov"`
+        (its per-step cost is already O(1) — see `HistoryEncoder`'s docstring),
+        or `AttentionHistory.init_cache()`'s real per-block KV cache under
+        `history="attention"`."""
+        return self.history_encoder.init_cache()
+
+    def history_step(self, token_feat: torch.Tensor, has_prev: torch.Tensor, cache) -> tuple[torch.Tensor, object]:
+        """One inference slot's worth of history encoding: advances `cache`
+        (from `init_history_cache`, or a previous `history_step` call) by
+        `token_feat`/`has_prev` (`(B, 1, ...)` — the just-emitted previous
+        token, same convention `giant.sample.sample_secondaries_ar` already
+        threads as `prev_repr`), and returns `(hist, new_cache)` — `hist` is
+        this slot's history summary (pass it as `_token_cond`'s `hist=` to
+        every model call made for this slot), `new_cache` is what to pass into
+        the *next* slot's `history_step`. Must be called exactly once per
+        slot — see `AttentionHistory.step`'s docstring."""
+        return self.history_encoder.step(token_feat, has_prev, cache)
+
+    def forward(
+        self,
+        x_t: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+        history_feat: torch.Tensor,
+        has_prev: torch.Tensor,
+        remaining_frac: torch.Tensor,
+        slot_idx: torch.Tensor,
+        t: torch.Tensor | None = None,
+        hist: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        B, K = x_t.shape[0], x_t.shape[1]
+        c_emb = self._token_cond(
+            cond_cont,
+            cond_cat,
+            stage1_out,
+            history_feat,
+            has_prev,
+            remaining_frac,
+            slot_idx,
+            hist=hist,
+        )
+        if self.time_emb is not None:
+            assert t is not None
+            t_emb = self.time_emb(t.reshape(-1)).view(B, K, -1)
+            cond = torch.cat([t_emb, c_emb], dim=-1)
+        else:
+            cond = c_emb
+        x_flat = x_t.reshape(B * K, -1)
+        cond_flat = cond.reshape(B * K, -1)
+        cond_cont_flat = cond_cont.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1)
+        cond_cat_flat = cond_cat.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1)
+        out = self.trunk(x_flat, cond_flat, cond_cont_flat, cond_cat_flat)
+        return out.view(B, K, -1)
+
+    def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
+        self._require_n_sec_head()
+        assert self.n_sec_head is not None
+        return self.n_sec_head(self._base_cond(cond_cont, cond_cat, stage1_out))
+
+    def predict_type(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+        history_feat: torch.Tensor,
+        has_prev: torch.Tensor,
+        remaining_frac: torch.Tensor,
+        slot_idx: torch.Tensor,
+        hist: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        self._require_type_head()
+        assert self.type_head is not None
+        c_emb = self._token_cond(
+            cond_cont,
+            cond_cat,
+            stage1_out,
+            history_feat,
+            has_prev,
+            remaining_frac,
+            slot_idx,
+            hist=hist,
+        )
+        B, K, _ = c_emb.shape
+        return self.type_head(c_emb.reshape(B * K, -1)).view(B, K, self.type_dim)
+
+    def predict_stop(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+        history_feat: torch.Tensor,
+        has_prev: torch.Tensor,
+        remaining_frac: torch.Tensor,
+        slot_idx: torch.Tensor,
+        hist: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """`(B, K)` raw stop logits — `n_sec.mode = "stop_token"` only.
+        Evaluated on slot `k`'s own conditioning (which carries slot `k-1`'s
+        history, same as `predict_type`), so this is `P(n_sec == k |
+        prefix)`: a high logit at slot `k` means "stop before generating a
+        token here" — the caller (`giant.sample.sample_secondaries_ar`)
+        checks it before spending a model call on that slot's token."""
+        self._require_stop_head()
+        assert self.stop_head is not None
+        c_emb = self._token_cond(
+            cond_cont,
+            cond_cat,
+            stage1_out,
+            history_feat,
+            has_prev,
+            remaining_frac,
+            slot_idx,
+            hist=hist,
+        )
+        B, K, _ = c_emb.shape
+        return self.stop_head(c_emb.reshape(B * K, -1)).view(B, K)
+
+
+class CriticModel(StageModel):
+    """Generator-agnostic WGAN-GP critic body: a scalar realism score, for
+    either stage (`stage="stage1"` mirrors v0.2 `Critic`; `stage="stage2"`
+    mirrors v0.2 `SecondaryCritic`, adding the same context-fusion path as
+    `Stage2OneShot`, via `StageModel._build_context_fusion`/`_cond_embed`).
+    Used only when that stage's `generator == "wgan"`.
+
+    Subclasses `StageModel` for the `cond_enc` construction and (stage 2)
+    context-fusion scaffolding only — its trunk is built directly via
+    `build_trunk` (output width 1) rather than through
+    `_build_trunk_and_heads`, since that helper is shaped around a
+    generator's `Objective`/time-embedding/flow-matching concerns
+    (`forward`'s `(x_t, cond) -> vector` shape) that don't apply to a critic's
+    `(x, cond) -> scalar` (gitea #57). `generator="wgan"` is passed to the
+    base purely because that's factually when a critic exists; nothing here
+    ever calls `_build_trunk_and_heads`, so no head/time-embedding machinery
+    is built from it. Never routed (MoE) — that's a separate, unrequested
+    axis of scope; see gitea #57's proposal, which covers only the trunk/
+    block registries."""
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        particle_cfg: ConditioningAxisConfig,
+        material_cfg: ConditioningAxisConfig,
+        in_dim: int,
+        hidden_dim: int = 256,
+        n_res_blocks: int = 6,
+        cond_out_dim: int = 128,
+        dropout: float = 0.0,
+        stage: str = "stage1",
+        context_dim: int = 64,
+        context_in_dim: int = X_DIM,
+        trunk_type: str = "resmlp",
+        block_conditioning: str = "add",
+    ) -> None:
+        super().__init__(
+            pdg_vocab,
+            mat_vocab,
+            particle_cfg,
+            material_cfg,
+            cond_out_dim=cond_out_dim,
+            generator="wgan",
+            noise_dim=0,
+        )
+        if stage not in ("stage1", "stage2"):
+            raise ValueError(f"stage must be 'stage1' or 'stage2', got {stage!r}")
+        self.stage = stage
+        if stage == "stage2":
+            self._build_context_fusion(context_in_dim, context_dim, cond_out_dim)
+        self.trunk = build_trunk(
+            None, trunk_type, in_dim, 1, hidden_dim, n_res_blocks, cond_out_dim, dropout, block_conditioning
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if self.stage == "stage2":
+            assert stage1_out is not None, "stage='stage2' CriticModel requires stage1_out"
+            cond = self._cond_embed(cond_cont, cond_cat, stage1_out)
+        else:
+            cond = self.cond_enc(cond_cont, cond_cat)
+        return self.trunk(x, cond, cond_cont, cond_cat).squeeze(-1)
@@ -0,0 +1,202 @@
+"""Generative objectives (flow/ddpm/wgan): `Objective` base + registry,
+mirroring `giant.model.routers`'s `Router` pattern (gitea #32). Each objective
+answers, in one place, the handful of questions every stage model/sampler/
+trainer used to re-derive independently from a bare `generator` string: does
+this stage need a time embedding, is it adversarial, does it fold the
+secondary type slice into its own trunk output, what does the trunk take as
+input, which stage-1/stage-2 loss does it train against.
+
+Self-contained (no dependency on `giant.model.models`, unlike `Router` which
+`giant.model.trunks` depends on) — `Objective` never needs to construct a
+stage model or critic itself, only describe one. This also sidesteps a
+`models.py` <-> `objectives.py` import cycle, since `models.py` calls
+`build_objective`.
+"""
+
+import inspect
+
+import torch
+
+from giant.model.schedule import (
+    CosineSchedule,
+    flow_matching_loss,
+    flow_matching_loss_secondary,
+    flow_matching_loss_secondary_ar,
+)
+
+# ---------------------------------------------------------------------------
+# Objective contract
+# ---------------------------------------------------------------------------
+
+
+class Objective:
+    """Contract for a pluggable generative objective. Not an `nn.Module` —
+    unlike `Router`, no objective owns learnable parameters, so a plain
+    strategy object is the honest fit.
+
+    `needs_time`/`is_adversarial`/`folds_type_slice`/`supports_stage2_decoder`
+    are set by each concrete subclass (no defaults here — a new objective
+    should have to state all four, not silently inherit one that happens to
+    be wrong for it). See `FlowObjective`/`DdpmObjective`/`WganObjective`.
+    """
+
+    needs_time: bool
+    is_adversarial: bool
+    folds_type_slice: bool
+    supports_stage2_decoder: bool = True
+
+    def trunk_in_dim(self, out_dim: int, noise_dim: int) -> int:
+        """Width of the trunk's own input — `out_dim` (denoising/flow-matching
+        a same-shape vector) for every non-adversarial objective;
+        `WganObjective` overrides to `noise_dim` (a single-pass noise-to-output
+        generator)."""
+        return out_dim
+
+    def build_schedule(self, n_steps: int, device: torch.device) -> CosineSchedule | None:
+        """Objective-owned auxiliary state a stage trainer must build once
+        and hold onto (device-placed) across its training loop. `None` for
+        every objective except `DdpmObjective` (its noise schedule)."""
+        return None
+
+    def stage1_loss(
+        self,
+        model: torch.nn.Module,
+        x1: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        *,
+        schedule: object | None = None,
+    ) -> torch.Tensor:
+        """Stage-1 training loss. Only implemented by non-adversarial
+        objectives — `WganObjective` is unused here, `WGANStageTrainer` has
+        its own G/D step instead."""
+        raise NotImplementedError(f"{type(self).__name__} has no stage1_loss")
+
+    def stage2_loss(
+        self,
+        model: torch.nn.Module,
+        x1_s2: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_ctx: torch.Tensor,
+        sec_mask: torch.Tensor,
+        *,
+        type_dim: int | None,
+        ar_inputs: dict[str, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        """Stage-2 secondary-decoder training loss, one-shot or
+        autoregressive depending on whether `ar_inputs` is given. Same
+        adversarial caveat as `stage1_loss`."""
+        raise NotImplementedError(f"{type(self).__name__} has no stage2_loss")
+
+
+OBJECTIVE_REGISTRY: dict[str, type[Objective]] = {}
+
+
+def register_objective(name: str):
+    def decorator(cls: type[Objective]) -> type[Objective]:
+        OBJECTIVE_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def build_objective(name: str, **kwargs) -> Objective:
+    """Factory: look up an `Objective` subclass by name (a `generator`
+    config value) from the registry.
+
+    Every registered objective is fed the same kwargs; kwargs not declared by
+    that type's constructor are silently dropped, so per-type hyperparameters
+    (e.g. `DdpmObjective`'s `n_steps`) can coexist in one call without
+    special-casing — same convention as `giant.model.routers.build_router`.
+    """
+    if name not in OBJECTIVE_REGISTRY:
+        raise ValueError(f"unknown generator/objective {name!r}; available: {sorted(OBJECTIVE_REGISTRY)}")
+    cls = OBJECTIVE_REGISTRY[name]
+    accepted = set(inspect.signature(cls.__init__).parameters) - {"self"}
+    filtered = {k: v for k, v in kwargs.items() if k in accepted}
+    return cls(**filtered)
+
+
+# ---------------------------------------------------------------------------
+# Concrete objectives
+# ---------------------------------------------------------------------------
+
+
+@register_objective("flow")
+class FlowObjective(Objective):
+    """Conditional flow matching (Lipman et al. 2022) — the primary
+    objective. ~10 ODE steps at inference (`giant.sample.sample_flow`)."""
+
+    needs_time = True
+    is_adversarial = False
+    folds_type_slice = False
+
+    def stage1_loss(self, model, x1, cond_cont, cond_cat, *, schedule=None) -> torch.Tensor:
+        return flow_matching_loss(model, x1, cond_cont, cond_cat)
+
+    def stage2_loss(
+        self,
+        model,
+        x1_s2,
+        cond_cont,
+        cond_cat,
+        stage1_ctx,
+        sec_mask,
+        *,
+        type_dim=None,
+        ar_inputs=None,
+    ) -> torch.Tensor:
+        if ar_inputs is not None:
+            return flow_matching_loss_secondary_ar(
+                model,
+                x1_s2,
+                cond_cont,
+                cond_cat,
+                stage1_ctx,
+                ar_inputs["history_feat"],
+                ar_inputs["has_prev"],
+                ar_inputs["remaining_frac"],
+                ar_inputs["slot_idx"],
+                sec_mask,
+                type_dim=type_dim,
+            )
+        return flow_matching_loss_secondary(model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=type_dim)
+
+
+@register_objective("ddpm")
+class DdpmObjective(Objective):
+    """Full DDPM ancestral sampling (Nichol & Dhariwal 2021 cosine schedule)
+    — the throwaway baseline. Stage-1 only: no `Stage2*` class has ever been
+    trained with `generator="ddpm"` in practice, so there's no stage-2 ddpm
+    loss to dispatch to (matches `FlowDDPMStageTrainer`'s pre-existing
+    stage-2 guard)."""
+
+    needs_time = True
+    is_adversarial = False
+    folds_type_slice = False
+    supports_stage2_decoder = False
+
+    def __init__(self, n_steps: int = 1000) -> None:
+        self.n_steps = n_steps
+
+    def build_schedule(self, n_steps: int, device: torch.device) -> CosineSchedule:
+        return CosineSchedule(T=n_steps).to(device)
+
+    def stage1_loss(self, model, x1, cond_cont, cond_cat, *, schedule=None) -> torch.Tensor:
+        assert schedule is not None, "DdpmObjective.stage1_loss needs a schedule (see build_schedule)"
+        return schedule.loss(model, x1, cond_cont, cond_cat)
+
+
+@register_objective("wgan")
+class WganObjective(Objective):
+    """WGAN-GP (Gulrajani et al. 2017) — single forward pass instead of an
+    ODE loop. `stage1_loss`/`stage2_loss` are unused: `WGANStageTrainer` owns
+    its own dual generator/critic step instead of a single scalar loss."""
+
+    needs_time = False
+    is_adversarial = True
+    folds_type_slice = True
+
+    def trunk_in_dim(self, out_dim: int, noise_dim: int) -> int:
+        return noise_dim
@@ -0,0 +1,414 @@
+"""Mixture-of-experts routing: `Router` base + registry, the four concrete
+router types, and composed/config-driven construction — self-contained, no
+dependency on any other `giant.model` submodule (issues.md Issue 8)."""
+
+import inspect
+import math
+import re
+from collections.abc import Sequence
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from giant.cond_layout import CondLayout
+from giant.constants import COND_DIM
+
+# ---------------------------------------------------------------------------
+# Routers — carried over unchanged from v0.2
+# ---------------------------------------------------------------------------
+
+
+class Router(nn.Module):
+    """Contract for a pluggable mixture-of-experts routing axis.
+
+    Subclasses implement `gate` (soft partition-of-unity weights over
+    experts, used in train mode for a fully differentiable mixture);
+    `top1` and `balance_loss` have working defaults so a new routing axis
+    is usually a one-method add. See `ROUTER_REGISTRY` / `build_router`.
+    """
+
+    def __init__(self, n_experts: int) -> None:
+        super().__init__()
+        self.n_experts = n_experts
+        self.gumbel = False
+        self.gumbel_tau = 1.0
+
+    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        """(B, n_experts) soft weights, rows summing to 1."""
+        raise NotImplementedError
+
+    def combine_weights(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        """(B, n_experts) train-time expert-combination weights.
+
+        Default (`gumbel=False`): identical to `gate()`. Opt-in
+        straight-through Gumbel-softmax (`gumbel=True`, train mode only):
+        hardens the forward pass to a one-hot sample (matching eval-time
+        top-1 dispatch) while keeping the soft sample's gradient on backward.
+
+        Forced fp32 (`torch.autocast(..., enabled=False)`) regardless of the
+        caller's ambient `train.precision` autocast region: `clamp_min(1e-8)`
+        below sits under bf16's precision but *above* fp16's ~6e-8 subnormal
+        floor, so `log_probs` degrading here is exactly the kind of quiet
+        drift that cost a whole rollout benchmark before (see the MoE section
+        of CLAUDE.md's Roadmap) — cheap to rule out (gitea #47).
+        """
+        with torch.autocast(cond_cont.device.type, enabled=False):
+            probs = self.gate(cond_cont, cond_cat)
+            if not (self.gumbel and self.training):
+                return probs
+            log_probs = torch.log(probs.clamp_min(1e-8))
+            return F.gumbel_softmax(log_probs, tau=self.gumbel_tau, hard=True, dim=-1)
+
+    def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        """(B,) hard expert index, used for eval-time grouped dispatch."""
+        with torch.autocast(cond_cont.device.type, enabled=False):
+            return self.gate(cond_cont, cond_cat).argmax(dim=-1)
+
+    def balance_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        """Importance CV^2 load-balancing auxiliary loss (Shazeer et al. 2017).
+
+        Forced fp32 — `importance` sums `gate()` over the whole batch (a
+        large-magnitude accumulation in reduced precision), then takes a
+        `std/mean` ratio: a classic catastrophic-cancellation shape (gitea
+        #47)."""
+        with torch.autocast(cond_cont.device.type, enabled=False):
+            importance = self.gate(cond_cont, cond_cat).sum(dim=0)  # (n_experts,)
+            return (importance.std() / (importance.mean() + 1e-8)) ** 2
+
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        """Optional supervised auxiliary loss shaping the router's own belief.
+
+        Default: none (a scalar 0). Routers gating on an unobservable
+        pre-step quantity (e.g. ProcessRouter) override this.
+        """
+        return torch.zeros((), device=cond_cont.device)
+
+    def entropy_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        """Optional auxiliary loss rewarding sharper (lower-entropy) routing."""
+        norm_entropy, _ = self.gate_stats(cond_cont, cond_cat)
+        return norm_entropy
+
+    def gate_stats(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Diagnostics: `(norm_entropy, importance)` — see v0.2 docstring for
+        the full explanation, unchanged in v0.3.0.
+
+        Forced fp32, same rationale as `balance_loss`/`combine_weights`: the
+        `+ 1e-8` epsilon here is `entropy_loss`'s training-loss path too, not
+        just a diagnostic (gitea #47)."""
+        with torch.autocast(cond_cont.device.type, enabled=False):
+            gate = self.gate(cond_cont, cond_cat)  # (B, n_experts)
+            row_entropy = -(gate * (gate + 1e-8).log()).sum(dim=-1)  # (B,)
+            norm_entropy = row_entropy.mean() / math.log(self.n_experts)
+            importance = gate.sum(dim=0)  # (n_experts,)
+            return norm_entropy, importance
+
+
+ROUTER_REGISTRY: dict[str, type[Router]] = {}
+
+
+def register_router(name: str):
+    def decorator(cls: type[Router]) -> type[Router]:
+        ROUTER_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def build_router(name: str, n_experts: int, **kwargs) -> Router:
+    """Factory: look up a `Router` subclass by name from the registry.
+
+    Every registered router type is fed the same `router` config dict;
+    kwargs not declared by that type's constructor are silently dropped, so
+    per-type hyperparameters (e.g. EnergyRouter's `temperature`) can coexist
+    in one config without special-casing.
+    """
+    if name not in ROUTER_REGISTRY:
+        raise ValueError(f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}")
+    cls = ROUTER_REGISTRY[name]
+    accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"}
+    filtered = {k: v for k, v in kwargs.items() if k in accepted}
+    return cls(n_experts=n_experts, **filtered)
+
+
+def _bounded_interp(raw: torch.Tensor, lo: float, hi: float) -> torch.Tensor:
+    """Sigmoid interpolation into `[lo, hi]` — smooth, always-positive-gradient
+    bound used for EnergyRouter's `learn_width`/`learn_temperature` modes."""
+    return lo + (hi - lo) * torch.sigmoid(raw)
+
+
+def _inverse_bounded_interp(value: float, lo: float, hi: float) -> float:
+    """Inverse of `_bounded_interp`, used once at construction to warm-start
+    `raw` so the initial effective width/temperature exactly equals `value`."""
+    p = min(max((value - lo) / (hi - lo), 1e-6), 1 - 1e-6)
+    return math.log(p / (1 - p))
+
+
+@register_router("none")
+class NoneRouter(Router):
+    """Uniform 1/n_experts gate — no learned routing signal at all.
+
+    Still builds n_experts expert trunks via RoutedTrunk (same parameter
+    budget as a real router), but every row gets an identical weight
+    regardless of conditioning. Ablates whether the *learned routing
+    signal* — as opposed to simply having multiple experts — is earning
+    its parameters. `top1()` (the base class default) always dispatches to
+    expert 0 (argmax of a uniform vector), which still exercises
+    RoutedTrunk's real per-expert grouped-dispatch code path at eval time.
+    """
+
+    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        B = cond_cont.shape[0]
+        return torch.full((B, self.n_experts), 1.0 / self.n_experts, device=cond_cont.device)
+
+
+@register_router("energy")
+class EnergyRouter(Router):
+    """Soft turn-on gate over normalized pre-step log-energy.
+
+    Reads `cond_cont[:, energy_idx]` (ignores cond_cat). `gate(e) =
+    softmax_i(-(e - c_i)^2 / tau)`; as tau -> 0 this hardens to
+    nearest-center (Voronoi) selection, exactly what `top1` uses at eval.
+    """
+
+    def __init__(
+        self,
+        n_experts: int = 4,
+        temperature: float = 0.5,
+        learn_centers: bool = True,
+        energy_idx: int = 3,
+        centers_init: Sequence[float] | None = None,
+        learn_width: bool = False,
+        learn_temperature: bool = False,
+        width_min_ratio: float = 0.1,
+        width_max_ratio: float = 10.0,
+    ) -> None:
+        super().__init__(n_experts)
+        if learn_width and learn_temperature:
+            raise ValueError("learn_width and learn_temperature are mutually exclusive")
+        self.temperature = temperature
+        self.energy_idx = energy_idx
+        self.learn_width = learn_width
+        self.learn_temperature = learn_temperature
+        if learn_width or learn_temperature:
+            if not (width_min_ratio < 1.0 < width_max_ratio):
+                raise ValueError(
+                    f"width_min_ratio ({width_min_ratio}) and width_max_ratio ({width_max_ratio}) must bracket 1.0"
+                )
+            self._width_lo = width_min_ratio * temperature
+            self._width_hi = width_max_ratio * temperature
+            raw0 = _inverse_bounded_interp(temperature, self._width_lo, self._width_hi)
+            if learn_width:
+                self.raw_width = nn.Parameter(torch.full((n_experts,), raw0))
+            else:
+                self.raw_temperature = nn.Parameter(torch.tensor(raw0))
+        if centers_init is None:
+            centers = torch.linspace(-2.0, 2.0, n_experts)
+        else:
+            if len(centers_init) != n_experts:
+                raise ValueError(f"centers_init has {len(centers_init)} values, expected n_experts={n_experts}")
+            centers = torch.tensor(list(centers_init), dtype=torch.float32)
+        if learn_centers:
+            self.centers = nn.Parameter(centers)
+        else:
+            self.register_buffer("centers", centers)
+
+    def effective_width(self) -> torch.Tensor | float:
+        if self.learn_width:
+            return _bounded_interp(self.raw_width, self._width_lo, self._width_hi)
+        if self.learn_temperature:
+            return _bounded_interp(self.raw_temperature, self._width_lo, self._width_hi)
+        return self.temperature
+
+    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        e = cond_cont[:, self.energy_idx].unsqueeze(-1)  # (B, 1)
+        d2 = (e - self.centers.unsqueeze(0)) ** 2  # (B, n_experts)
+        return torch.softmax(-d2 / self.effective_width(), dim=-1)
+
+
+@register_router("pdg")
+class PdgRouter(Router):
+    """Soft turn-on gate over a learned PDG embedding (own table, separate
+    from the trunk's `ConditionEncoder`). No supervision needed — PDG code
+    is already known at pre-step time."""
+
+    def __init__(
+        self,
+        n_experts: int,
+        pdg_vocab: int,
+        emb_dim: int = 8,
+        temperature: float = 0.5,
+        learn_centers: bool = True,
+    ) -> None:
+        super().__init__(n_experts)
+        self.temperature = temperature
+        self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
+        centers = torch.randn(n_experts, emb_dim) * 0.1
+        if learn_centers:
+            self.centers = nn.Parameter(centers)
+        else:
+            self.register_buffer("centers", centers)
+
+    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        e = self.pdg_emb(cond_cat[:, CondLayout.PDG_COL])  # (B, emb_dim)
+        d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(-1)  # (B, n_experts)
+        return torch.softmax(-d2 / self.temperature, dim=-1)
+
+
+@register_router("process")
+class ProcessRouter(Router):
+    """Routes on the physics process expected to end the step — a post-step
+    outcome, so a small classifier over pre-step conditioning predicts it
+    (own pdg/material embeddings, separate from the trunk's ConditionEncoder).
+    `n_experts` doubles as the number of process classes. Supervised via
+    `classify_loss` against the true `process` label at train time only;
+    `gate`/`top1` never see it."""
+
+    def __init__(
+        self,
+        n_experts: int,
+        pdg_vocab: int,
+        mat_vocab: int,
+        emb_dim: int = 8,
+        hidden_dim: int = 64,
+    ) -> None:
+        super().__init__(n_experts)
+        self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
+        self.mat_emb = nn.Embedding(mat_vocab, emb_dim)
+        self.classifier = nn.Sequential(
+            nn.Linear(COND_DIM + 2 * emb_dim, hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, n_experts),
+        )
+
+    def logits(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        pdg_e = self.pdg_emb(cond_cat[:, CondLayout.PDG_COL])
+        mat_e = self.mat_emb(cond_cat[:, CondLayout.MAT_COL])
+        h = torch.cat([cond_cont, pdg_e, mat_e], dim=-1)
+        return self.classifier(h)
+
+    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        return torch.softmax(self.logits(cond_cont, cond_cat), dim=-1)
+
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        return F.cross_entropy(self.logits(cond_cont, cond_cat), labels)
+
+
+class ComposedRouter(Router):
+    """Joint router over independent axes (e.g. energy x pdg), outer-product
+    gated. Not registered in `ROUTER_REGISTRY`; use `build_composed_router`."""
+
+    def __init__(self, routers: list[Router]) -> None:
+        if not routers:
+            raise ValueError("ComposedRouter needs at least one sub-router")
+        n_experts = 1
+        for r in routers:
+            n_experts *= r.n_experts
+        super().__init__(n_experts)
+        self.routers = nn.ModuleList(routers)
+
+    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        joint = self.routers[0].gate(cond_cont, cond_cat)  # (B, n_0)
+        for router in self.routers[1:]:
+            g = router.gate(cond_cont, cond_cat)  # (B, n_i)
+            joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(1)  # (B, prod so far)
+        return joint
+
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        total = torch.zeros((), device=cond_cont.device)
+        for router in self.routers:
+            total = total + router.classify_loss(cond_cont, cond_cat, labels)
+        return total
+
+
+def build_composed_router(specs: list[dict], **shared_kwargs) -> ComposedRouter:
+    """Build a `ComposedRouter` from a list of per-axis router specs — see
+    `_parse_composed_axes`."""
+    routers = [
+        build_router(
+            spec["type"],
+            spec["n_experts"],
+            **{
+                **shared_kwargs,
+                **{k: v for k, v in spec.items() if k not in ("type", "n_experts")},
+            },
+        )
+        for spec in specs
+    ]
+    return ComposedRouter(routers)
+
+
+_AXIS_KEY_RE = re.compile(r"^axis(\d+)_(.+)$")
+
+
+def _parse_composed_axes(router_cfg: dict) -> list[dict]:
+    """Regroup `axis{i}_{field}` flat keys into a list of per-axis spec dicts.
+
+    e.g. `axis0_type = "energy"`, `axis0_n_experts = 4`, `axis1_type = "pdg"`,
+    `axis1_n_experts = 3`, `axis1_emb_dim = 8`. Axis indices must be
+    contiguous from 0.
+    """
+    axes: dict[int, dict] = {}
+    for key, value in router_cfg.items():
+        m = _AXIS_KEY_RE.match(key)
+        if m is None:
+            continue
+        idx, field = int(m.group(1)), m.group(2)
+        axes.setdefault(idx, {})[field] = value
+    missing = set(range(len(axes))) - axes.keys()
+    if missing:
+        raise ValueError(f"composed router config has gaps at axis indices {missing}")
+    return [axes[i] for i in range(len(axes))]
+
+
+# Router types that read cond_cat's pdg index through their own
+# nn.Embedding(pdg_vocab, ...), regardless of the trunk's particle
+# conditioning mode — see _check_router_conditioning_compat.
+_VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process")
+
+
+def _check_router_conditioning_compat(router_types: list[str], particle_conditioning: str) -> None:
+    """Reject a router axis that reintroduces a training-vocab PDG lookup
+    under `conditioning.particle.type = "physical"`.
+
+    `PdgRouter`/`ProcessRouter` always build their own dataset-scoped
+    `nn.Embedding(pdg_vocab, ...)`, independent of `ConditionEncoder`'s
+    particle mode. Pairing either with `"physical"` would silently
+    reintroduce a training-menu-scoped lookup at the routing layer,
+    defeating the point of physical-property conditioning. Raised loudly at
+    model-build time.
+    """
+    bad = sorted(set(router_types) & set(_VOCAB_SCOPED_ROUTER_TYPES))
+    if bad and particle_conditioning == "physical":
+        raise ValueError(
+            f"router type(s) {bad} always use a training-vocab PDG embedding, "
+            "which is incompatible with conditioning.particle.type='physical' "
+            "(whose whole point is generalizing beyond that vocab) — pick a "
+            "different router type (e.g. 'energy') or use "
+            "conditioning.particle.type='embedding'."
+        )
+
+
+def _build_router_from_cfg(
+    router_cfg: dict,
+    pdg_vocab: int,
+    mat_vocab: int,
+    particle_conditioning: str = "embedding",
+) -> Router:
+    """Resolve one stage's `router` config into a `Router`, single-axis or
+    composed. `gumbel` is set as a post-construction attribute (shared by
+    every router type, not a per-type constructor kwarg)."""
+    shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
+    if router_cfg["type"] == "composed":
+        axes = _parse_composed_axes(router_cfg)
+        _check_router_conditioning_compat([a["type"] for a in axes], particle_conditioning)
+        router = build_composed_router(axes, **shared_vocab)
+        router.gumbel = bool(router_cfg.get("gumbel", False))
+        return router
+    _check_router_conditioning_compat([router_cfg["type"]], particle_conditioning)
+    router_kwargs = {k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")}
+    router_kwargs.setdefault("pdg_vocab", pdg_vocab)
+    router_kwargs.setdefault("mat_vocab", mat_vocab)
+    router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
+    router.gumbel = bool(router_cfg.get("gumbel", False))
+    return router
@@ -11,9 +11,7 @@ class CosineSchedule:
        steps = np.arange(T + 1, dtype=np.float64)
        f = np.cos(((steps / T + s) / (1.0 + s)) * np.pi / 2.0) ** 2
        alpha_bars = (f / f[0]).astype(np.float32)
-        betas = np.clip(1.0 - alpha_bars[1:] / alpha_bars[:-1], 0.0, 0.999).astype(
-            np.float32
-        )
+        betas = np.clip(1.0 - alpha_bars[1:] / alpha_bars[:-1], 0.0, 0.999).astype(np.float32)

        self.betas = torch.from_numpy(betas)
        self.alphas = torch.from_numpy((1.0 - betas))
@@ -90,8 +88,8 @@ def flow_matching_loss_secondary(
        "physical"`'s width and the only case this function handled before
        v0.3.0 step 4. `0` means no type slice is in `x1` at all (`target`
        in `("onehot", "embedding")` under `generator in ("flow", "ddpm")` —
-        see docs/v0.3.0-design.md decision 2, `Stage2OneShot.type_head`
-        handles the type loss separately in that case).
+        `Stage2OneShot.type_head` handles the type loss separately in that
+        case).

    Only valid-slot dimensions contribute to the loss; padded slots are zeroed
    before averaging, so the loss is not diluted by empty slots.
@@ -156,10 +154,10 @@ def flow_matching_loss_secondary_ar(
    model call signature differs enough (four extra per-token conditioning
    tensors) that merging would need an awkward shape-flag + closure.

-    Under teacher forcing (docs/v0.3.0-design.md §6.2 point 3) this is still a
+    Under teacher forcing this is still a
    single parallel pass over all K_MAX tokens — `x1`/`history_feat`/etc. are
    already built from ground truth for every slot by the caller
-    (`giant.train._assemble_stage2_ar_inputs`/`_assemble_stage2_ar_target`).
+    (`giant.training.stage2_inputs._assemble_stage2_ar_inputs`/`_assemble_stage2_ar_target`).

    x1: (B, K_MAX, CONT_SLOT_DIM + type_dim) — per-token flattened target
        (stick_logit, dir, then a `type_dim`-wide type slice)
@@ -0,0 +1,318 @@
+"""Build-only model introspection (gitea #46): construct the resolved
+Stage1/Stage2/critic graph from a config with no dataset attached, and report
+per-module parameter counts, trunk widths, which heads exist, and — via
+differential probing — which `conditioning`/`stage1_model`/`stage2_model`
+config keys actually shape the built model. This is the runtime counterpart
+to `tests/test_config_consumed_keys.py`'s static per-identifier audit: that
+test asks "does any code reference this key's name at all", this module asks
+"given *this* resolved config, does the key change what `build_models`/
+`build_critics` (`giant/model/builders.py`) actually produces".
+
+Differential probing, not identifier matching: build the model once from the
+resolved config and take a structural fingerprint (`_fingerprint` — which
+submodules exist, every parameter's/buffer's shape+dtype, every plain scalar
+attribute stored on any module). Then, for each in-scope leaf key, perturb
+just that one value (`_perturb`), rebuild, and re-fingerprint. A changed
+fingerprint — or a rebuild that raises — means the key was consumed; an
+identical fingerprint means construction never looked at it under this
+particular config. A key can be genuinely inert under one config and live
+under another (e.g. any `stage1_model.router.*` key when `router.enabled =
+false`) — that config-dependence is exactly the "silently degenerate
+combination" issue #46 is after, so it is reported per-run rather than
+baked into a static table.
+
+Keys legitimately owned by the trainer/sampler/rollout rather than by
+`build_models`/`build_critics` (loss weights, WGAN-GP training
+hyperparameters, teacher-forcing and stage1-context schedules, ...) are
+cataloged in `_NOT_BUILD_TIME` below so the report doesn't flag them as
+suspicious. One leaf is inert under every config today —
+`stage2_model.autoregressive.order` — matching
+`tests/test_config_consumed_keys.py`'s own `_KNOWN_UNUSED` entry; it is
+deliberately *not* in `_NOT_BUILD_TIME`, since "always inert" is itself the
+finding those two tests independently converge on.
+"""
+
+import copy
+from dataclasses import dataclass, field
+
+import torch.nn as nn
+
+from giant.config import _get_path, _set_path, leaf_paths
+from giant.model.builders import build_critics, build_models
+from giant.model.trunks import RoutedTrunk
+
+_IN_SCOPE_ROOTS = ("conditioning", "stage1_model", "stage2_model")
+
+_PROBE_STR = "__giant_model_summary_probe__"
+
+# A handful of string leaves branch on equality against one specific literal
+# (e.g. `builders.py`: `stop_token = s2_spec.n_sec.mode == "stop_token"`),
+# where every value other than that literal behaves identically. A single
+# generic sentinel probe would then falsely read as inert whenever the
+# config's *current* value is already one of those identically-behaving
+# "other" values (e.g. mode="head") — it never crosses the one boundary that
+# actually matters. Named here so probing tries the real alternative(s) too;
+# every other string leaf is registry-validated (raises on garbage, still
+# correctly detected as consumed) or genuinely value-independent, so doesn't
+# need an entry.
+_STRING_ALTERNATIVES: dict[str, tuple[str, ...]] = {
+    "stage2_model.n_sec.owner": ("stage1", "stage2"),
+    "stage2_model.n_sec.mode": ("stop_token", "head", "truth"),
+    "stage2_model.particle_type.target": ("physical", "onehot", "embedding"),
+}
+
+# Verified by reading giant/training/trainers.py, giant/training/stage2_inputs.py
+# and giant/rollout.py while implementing gitea #46 — not auto-derived, so a
+# future reader touching these fields should re-check this table still holds.
+_NOT_BUILD_TIME: dict[str, str] = {
+    "stage1_model.init_from": "training/checkpoint.py's init_stages_from_checkpoints, run before build_stage_trainers (gitea #42)",
+    "stage1_model.freeze": "trainers.py: StageSpec.freeze, gates StageTrainer._step_optimizer (gitea #42)",
+    "stage2_model.init_from": "training/checkpoint.py's init_stages_from_checkpoints, run before build_stage_trainers (gitea #42)",
+    "stage2_model.freeze": "trainers.py: StageSpec.freeze, gates StageTrainer._step_optimizer (gitea #42)",
+    "stage1_model.lambda": "trainers.py: StageSpec.lambda_weight, the total-loss mix weight",
+    "stage2_model.lambda": "trainers.py: StageSpec.lambda_weight, the total-loss mix weight",
+    "stage2_model.n_sec.lambda": "trainers.py: StageSpec.n_sec_lambda, the n_sec-head loss weight",
+    "stage2_model.particle_type.lambda": "trainers.py: Stage2Trainer.particle_type_lambda, the type-head loss weight",
+    "stage2_model.particle_type.other_policy": "giant/rollout.py: resolves an 'other'-bucket secondary's PDG code at inference",
+    "stage2_model.particle_type.class_weighting": "trainers.py: FlowDDPMStageTrainer.type_class_weights, shapes the type-head loss, not the built graph (gitea #44)",
+    "stage2_model.autoregressive.teacher_forcing": "giant/training/stage2_inputs.py's training-time input assembly",
+    "stage2_model.autoregressive.tf_p_start": "trainers.py's teacher-forcing schedule",
+    "stage2_model.autoregressive.tf_p_end": "trainers.py's teacher-forcing schedule",
+    "stage2_model.stage1_context": "trainers.py's stage1/stage2 boundary — StageTrainer._stage1_context",
+    "stage2_model.ctx_p_start": "trainers.py's stage1-context sampling schedule",
+    "stage2_model.ctx_p_end": "trainers.py's stage1-context sampling schedule",
+    "stage1_model.router.lambda_balance": "trainers.py's load-balancing auxiliary loss weight",
+    "stage1_model.router.lambda_entropy": "trainers.py's entropy-regularization auxiliary loss weight",
+    "stage1_model.router.lambda_proc": "trainers.py's supervised process-classification auxiliary loss weight",
+    "stage1_model.router.gumbel_tau_start": "trainers.py's expert-combination Gumbel-softmax temperature anneal",
+    "stage1_model.router.gumbel_tau_end": "trainers.py's expert-combination Gumbel-softmax temperature anneal",
+    "stage2_model.router.lambda_balance": "trainers.py's load-balancing auxiliary loss weight",
+    "stage2_model.router.lambda_entropy": "trainers.py's entropy-regularization auxiliary loss weight",
+    "stage2_model.router.lambda_proc": "trainers.py's supervised process-classification auxiliary loss weight",
+    "stage2_model.router.gumbel_tau_start": "trainers.py's expert-combination Gumbel-softmax temperature anneal",
+    "stage2_model.router.gumbel_tau_end": "trainers.py's expert-combination Gumbel-softmax temperature anneal",
+    "stage1_model.wgan.n_critic": "trainers.py's WGAN-GP critic-update cadence",
+    "stage1_model.wgan.gp_weight": "trainers.py's WGAN-GP gradient-penalty coefficient",
+    "stage1_model.wgan.critic_lr": "trainers.py's critic optimizer learning rate",
+    "stage2_model.wgan.n_critic": "trainers.py's WGAN-GP critic-update cadence",
+    "stage2_model.wgan.gp_weight": "trainers.py's WGAN-GP gradient-penalty coefficient",
+    "stage2_model.wgan.critic_lr": "trainers.py's critic optimizer learning rate",
+    "stage2_model.wgan.gumbel_tau_start": "trainers.py's type-slice Gumbel-softmax temperature anneal (type_gumbel_tau_start)",
+    "stage2_model.wgan.gumbel_tau_end": "trainers.py's type-slice Gumbel-softmax temperature anneal (type_gumbel_tau_end)",
+}
+
+
+@dataclass
+class ModelSummary:
+    modules: dict[str, nn.Module]
+    consumed: list[str]
+    inert: list[str]
+    elsewhere: list[str]
+    pdg_vocab: int
+    mat_vocab: int
+    vocab_caveats: list[str] = field(default_factory=list)
+
+
+def _build_model_config(cfg: dict, pdg_vocab: int, mat_vocab: int) -> dict:
+    return {
+        "pdg_vocab": pdg_vocab,
+        "mat_vocab": mat_vocab,
+        "conditioning": cfg["conditioning"],
+        "stage1_model": cfg["stage1_model"],
+        "stage2_model": cfg["stage2_model"],
+    }
+
+
+def _built_modules(cfg: dict, pdg_vocab: int, mat_vocab: int) -> dict[str, nn.Module]:
+    model_config = _build_model_config(cfg, pdg_vocab, mat_vocab)
+    modules: dict[str, nn.Module] = {}
+    for name, m in build_models(model_config).items():
+        if m is not None:
+            modules[name] = m
+    for name, m in build_critics(model_config).items():
+        if m is not None:
+            modules[f"{name}_critic"] = m
+    return modules
+
+
+def _fingerprint(modules: dict[str, nn.Module]) -> list:
+    """A config-shape fingerprint of the built graph: which submodules
+    exist, every parameter's/buffer's shape+dtype (never values — those are
+    randomly initialized and irrelevant to *structure*), and every plain
+    scalar attribute any module stores on itself (e.g. `Stage2Autoregressive
+    .stop_sampling`, `EnergyRouter.temperature`) — this is what makes a
+    non-parametric key's effect on construction observable."""
+    sig = []
+    for stage_name, module in modules.items():
+        for mod_name, m in module.named_modules():
+            full = f"{stage_name}.{mod_name}" if mod_name else stage_name
+            for k, v in vars(m).items():
+                if k.startswith("_"):
+                    continue
+                if v is None or isinstance(v, (bool, int, float, str)):
+                    sig.append((full, k, v))
+        for pname, p in module.named_parameters():
+            sig.append((stage_name, "param", pname, tuple(p.shape), str(p.dtype)))
+        for bname, b in module.named_buffers():
+            sig.append((stage_name, "buffer", bname, tuple(b.shape), str(b.dtype)))
+    return sorted(sig, key=repr)
+
+
+def _perturb_candidates(path: str, value) -> list:
+    """Values to try perturbing `path`'s current `value` to, in order —
+    probing stops at the first one that changes the fingerprint or raises.
+    Almost always a single candidate; see `_STRING_ALTERNATIVES`."""
+    if isinstance(value, bool):
+        return [not value]
+    if isinstance(value, int):
+        return [value + 1]
+    if isinstance(value, float):
+        return [value + 1.0]
+    if isinstance(value, str):
+        alternatives = [v for v in _STRING_ALTERNATIVES.get(path, ()) if v != value]
+        return [*alternatives, _PROBE_STR]
+    raise TypeError(f"gitea #46 probing: unsupported leaf value type {type(value)!r} ({value!r})")
+
+
+def _vocab_caveats(cfg: dict) -> list[str]:
+    caveats = []
+    if _get_path(cfg, "conditioning.particle.type") == "embedding":
+        caveats.append(
+            "conditioning.particle.type = 'embedding' -- pdg_vocab below is a "
+            "placeholder (no dataset attached to derive the real training vocab size)"
+        )
+    if _get_path(cfg, "conditioning.material.type") == "embedding":
+        caveats.append(
+            "conditioning.material.type = 'embedding' -- mat_vocab below is a "
+            "placeholder (no dataset attached to derive the real training vocab size)"
+        )
+    for stage in ("stage1_model", "stage2_model"):
+        router_type = _get_path(cfg, f"{stage}.router.type")
+        if _get_path(cfg, f"{stage}.router.enabled") and router_type in ("pdg", "process"):
+            caveats.append(
+                f"{stage}.router.type = {router_type!r} builds its own pdg_vocab-sized "
+                "embedding -- the count above is a placeholder"
+            )
+    return caveats
+
+
+def summarize_model(cfg: dict, pdg_vocab: int, mat_vocab: int) -> ModelSummary:
+    """Build `cfg`'s model with no dataset attached and report its resolved
+    graph, plus which `conditioning`/`stage1_model`/`stage2_model` config
+    keys actually shaped it (differential probing — see module docstring).
+    `cfg` must already be a fully-merged v0.3 config (`merge_cli_overrides`
+    output) — this does not migrate or validate it."""
+    modules = _built_modules(cfg, pdg_vocab, mat_vocab)
+    baseline_fp = _fingerprint(modules)
+
+    in_scope = [p for p in leaf_paths(cfg) if p.split(".", 1)[0] in _IN_SCOPE_ROOTS]
+    consumed: list[str] = []
+    inert: list[str] = []
+    elsewhere: list[str] = []
+    for path in in_scope:
+        original = _get_path(cfg, path)
+        changed = False
+        for candidate in _perturb_candidates(path, original):
+            probe_cfg = copy.deepcopy(
+                {
+                    "conditioning": cfg["conditioning"],
+                    "stage1_model": cfg["stage1_model"],
+                    "stage2_model": cfg["stage2_model"],
+                }
+            )
+            _set_path(probe_cfg, path, candidate)
+            try:
+                changed = _fingerprint(_built_modules(probe_cfg, pdg_vocab, mat_vocab)) != baseline_fp
+            except Exception:
+                changed = True
+            if changed:
+                break
+        if changed:
+            consumed.append(path)
+        elif path in _NOT_BUILD_TIME:
+            elsewhere.append(path)
+        else:
+            inert.append(path)
+
+    return ModelSummary(
+        modules=modules,
+        consumed=sorted(consumed),
+        inert=sorted(inert),
+        elsewhere=sorted(elsewhere),
+        pdg_vocab=pdg_vocab,
+        mat_vocab=mat_vocab,
+        vocab_caveats=_vocab_caveats(cfg),
+    )
+
+
+def _tree_lines(module: nn.Module, name: str, indent: int = 0) -> list[str]:
+    total = sum(p.numel() for p in module.parameters())
+    in_dim = getattr(module, "in_dim", None)
+    out_dim = getattr(module, "out_dim", None)
+    widths = f"  [in={in_dim}, out={out_dim}]" if in_dim is not None and out_dim is not None else ""
+    lines = [f"{'  ' * indent}{name} ({type(module).__name__}): {total:,}{widths}"]
+    for child_name, child in module.named_children():
+        lines.extend(_tree_lines(child, child_name, indent + 1))
+    return lines
+
+
+_HEAD_NAMES = ("n_sec_head", "type_head", "stop_head")
+
+
+def _stage_header(name: str, module: nn.Module) -> list[str]:
+    total = sum(p.numel() for p in module.parameters())
+    lines = [f"{name}: {type(module).__name__} -- {total:,} parameters"]
+    generator = getattr(module, "generator_kind", None)
+    if generator is not None:
+        lines.append(f"  generator: {generator}")
+    trunk = getattr(module, "trunk", None)
+    if trunk is not None:
+        in_dim = getattr(trunk, "in_dim", "?")
+        out_dim = getattr(trunk, "out_dim", "?")
+        if isinstance(trunk, RoutedTrunk):
+            detail = f"routed, n_experts={trunk.router.n_experts}, expert type={type(trunk.experts[0]).__name__}"
+        else:
+            detail = f"unrouted, {type(trunk).__name__}"
+        lines.append(f"  trunk: {detail}, in={in_dim}, out={out_dim}")
+    history_kind = getattr(module, "history_kind", None)
+    if history_kind is not None:
+        lines.append(f"  autoregressive history: {history_kind}")
+    present = [h for h in _HEAD_NAMES if getattr(module, h, None) is not None]
+    absent = [h for h in _HEAD_NAMES if hasattr(module, h) and getattr(module, h) is None]
+    if present or absent:
+        lines.append(f"  heads present: {', '.join(present) if present else 'none'}")
+        if absent:
+            lines.append(f"  heads absent: {', '.join(absent)}")
+    return lines
+
+
+def render_summary(summary: ModelSummary) -> str:
+    lines: list[str] = []
+    for name, module in summary.modules.items():
+        lines.extend(_stage_header(name, module))
+        lines.extend(_tree_lines(module, name, indent=1))
+        lines.append("")
+
+    lines.append(
+        f"config keys read during construction: {len(summary.consumed)} / "
+        f"read elsewhere (trainer/sampler/rollout): {len(summary.elsewhere)} / "
+        f"inert under this config: {len(summary.inert)}"
+    )
+    if summary.elsewhere:
+        lines.append("read elsewhere, not by construction:")
+        for path in summary.elsewhere:
+            lines.append(f"  {path}  ({_NOT_BUILD_TIME[path]})")
+    lines.append("inert under this config (declared, parsed, but doing nothing here):")
+    if summary.inert:
+        for path in summary.inert:
+            lines.append(f"  {path}")
+    else:
+        lines.append("  (none)")
+
+    if summary.vocab_caveats:
+        lines.append("")
+        lines.append("vocab placeholder caveats:")
+        for caveat in summary.vocab_caveats:
+            lines.append(f"  {caveat}")
+
+    return "\n".join(lines)
@@ -0,0 +1,292 @@
+"""Trunks: everything downstream of the fused conditioning vector — a
+registrable expert *body* architecture (`TRUNK_REGISTRY`/`register_trunk`),
+used standalone or mixed by a `Router` (issues.md Issue 8; trunk-selectability
+gitea #33).
+
+Whether a body is mixed is orthogonal to which body it is: `RoutedTrunk`
+builds `router.n_experts` instances of whichever body `trunk_type` names, so
+a future body (e.g. a transformer) automatically gets a mixture variant for
+free — no separate "routed transformer trunk" class needed.
+"""
+
+import torch
+import torch.nn as nn
+
+from giant.model.layers import build_block
+from giant.model.routers import Router
+
+TRUNK_REGISTRY: dict[str, type[nn.Module]] = {}
+
+
+def register_trunk(name: str):
+    def decorator(cls: type[nn.Module]) -> type[nn.Module]:
+        TRUNK_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def build_expert_body(
+    name: str,
+    in_dim: int,
+    out_dim: int,
+    hidden_dim: int,
+    n_blocks: int,
+    cond_dim: int,
+    dropout: float = 0.0,
+    block_conditioning: str = "add",
+) -> nn.Module:
+    """Factory: look up a registered trunk body by name and construct one
+    instance of it — used both for a standalone (unrouted) trunk and for each
+    expert inside a `RoutedTrunk`. `block_conditioning` selects the
+    `BLOCK_REGISTRY` entry each body's internal `ResBlock`-family blocks use
+    (`trunk.block_conditioning`, gitea #34) — an optional trailing kwarg a
+    future non-`ResBlock`-based body can simply ignore, same idiom as
+    `Trunk.forward`'s accept-and-ignore `cond_cont`/`cond_cat`."""
+    if name not in TRUNK_REGISTRY:
+        raise ValueError(f"unknown trunk type {name!r}; available: {sorted(TRUNK_REGISTRY)}")
+    cls = TRUNK_REGISTRY[name]
+    return cls(in_dim, out_dim, hidden_dim, n_blocks, cond_dim, dropout, block_conditioning=block_conditioning)
+
+
+@register_trunk("resmlp")
+class ExpertTrunk(nn.Module):
+    """`input_proj -> ResBlock stack -> out_proj` — the registered `"resmlp"`
+    trunk body. Used both standalone (no router: `forward`'s `cond_cont`/
+    `cond_cat` are accepted and ignored, satisfying the `Trunk` interface
+    directly with no wrapper class) and as one expert inside a `RoutedTrunk`
+    (`_route_forward` calls it with just `(x, cond)`).
+
+    Unlike v0.2, `out_dim` is independent of `in_dim` — needed by stage-2 AR
+    tokens later (`noise_dim` in, `4 + type_dim` out), even though every
+    step-2/3 caller still has `in_dim == out_dim`.
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        hidden_dim: int,
+        n_blocks: int,
+        cond_dim: int,
+        dropout: float = 0.0,
+        block_conditioning: str = "add",
+    ) -> None:
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.input_proj = nn.Linear(in_dim, hidden_dim)
+        self.blocks = nn.ModuleList(
+            [build_block(block_conditioning, hidden_dim, cond_dim, dropout) for _ in range(n_blocks)]
+        )
+        self.out_proj = nn.Linear(hidden_dim, out_dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond: torch.Tensor,
+        cond_cont: torch.Tensor | None = None,
+        cond_cat: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        x = self.input_proj(x)
+        for block in self.blocks:
+            x = block(x, cond)
+        return self.out_proj(x)
+
+
+@register_trunk("linear")
+class LinearTrunk(nn.Module):
+    """`nn.Linear(in_dim + cond_dim, out_dim)` over `concat([x, cond])` —
+    the trivial trunk body: no hidden layer, no ResBlock stack, no
+    nonlinearity. Ablates whether trunk depth/nonlinearity is earning its
+    parameters, holding everything else (heads, ConditionEncoder,
+    generator, ...) fixed. Composes for free with `router.enabled = true`
+    (gitea #33): a RoutedTrunk of n_experts linear bodies is "mixture of
+    trivial linear experts". `hidden_dim`/`n_blocks`/`dropout`/
+    `block_conditioning` are accepted and ignored, matching
+    `build_expert_body`'s shared factory signature.
+
+    `x` — the trunk's own input (e.g. the noised primary vector for flow
+    matching) — does not already carry conditioning; that's fused in
+    per-body via `cond`. So this concatenates `x` and `cond` itself to
+    remain a valid, conditioning-dependent model.
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        hidden_dim: int,
+        n_blocks: int,
+        cond_dim: int,
+        dropout: float = 0.0,
+        block_conditioning: str = "add",
+    ) -> None:
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.linear = nn.Linear(in_dim + cond_dim, out_dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond: torch.Tensor,
+        cond_cont: torch.Tensor | None = None,
+        cond_cat: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.linear(torch.cat([x, cond], dim=-1))
+
+
+def _route_forward(
+    experts: nn.ModuleList,
+    router: Router,
+    x: torch.Tensor,
+    cond: torch.Tensor,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+    training: bool,
+) -> torch.Tensor:
+    """Shared dispatch for `RoutedTrunk`.
+
+    Train mode: full mixture `sum_i weight_i * expert_i(x)` — always
+    N-expert dense compute, fully differentiable (`weight` is
+    `router.combine_weights`). Eval mode: grouped top-1 dispatch — each row
+    runs exactly one expert, the actual source of the per-call speedup.
+
+    The accumulator's dtype is deferred to the first expert call rather than
+    fixed at fp32: under autocast (`train.precision = "bf16"`, gitea #47) an
+    expert's `ResBlock` stack returns bf16, and an fp32-fixed accumulator
+    would silently upcast every mixture term (train mode) or downcast every
+    dispatched row via `index_put_` (eval mode) — making a `RoutedTrunk`
+    return a different dtype than the unrouted `ExpertTrunk` it's a drop-in
+    replacement for, purely because `router.enabled` was set.
+
+    `router.combine_weights` is deliberately fp32 internally (it forces its
+    own autocast-disabled region — see `Router.combine_weights`'s docstring),
+    so `weights` itself is always fp32 regardless of the ambient precision.
+    Left as-is, `weights[:, i:i+1] * expert(x, cond)` would type-promote the
+    whole mixture back to fp32 by ordinary PyTorch promotion rules — the same
+    dtype-mismatch bug this function exists to avoid, just moved one line
+    over. `weights` is cast down to each expert's own output dtype right
+    before combining: the softmax stays numerically stable at fp32, but its
+    *result* (values in [0, 1], not precision-sensitive to represent) loses
+    nothing meaningful by then being used at bf16.
+    """
+    if training:
+        weights = router.combine_weights(cond_cont, cond_cat)  # (B, n_experts), fp32
+        out = None
+        for i, expert in enumerate(experts):
+            expert_out = expert(x, cond)
+            term = weights[:, i : i + 1].to(expert_out.dtype) * expert_out
+            out = term if out is None else out + term
+        assert out is not None, "RoutedTrunk built with zero experts"
+        return out
+
+    idx = router.top1(cond_cont, cond_cat)  # (B,)
+    out = None
+    for i, expert in enumerate(experts):
+        mask = idx == i
+        if mask.any():
+            expert_out = expert(x[mask], cond[mask])
+            if out is None:
+                out = torch.zeros(x.shape[0], expert_out.shape[-1], device=x.device, dtype=expert_out.dtype)
+            out[mask] = expert_out
+    if out is None:
+        # No row was ever dispatched (only reachable with an empty batch,
+        # x.shape[0] == 0) — nothing to infer a dtype from, so fall back to
+        # x's own, matching this function's pre-autocast behavior.
+        out = torch.zeros(x.shape[0], experts[0].out_dim, device=x.device, dtype=x.dtype)
+    return out
+
+
+class Trunk(nn.Module):
+    """Interface implemented by a standalone trunk body (any `TRUNK_REGISTRY`
+    entry, e.g. `ExpertTrunk`) and by `RoutedTrunk`: everything downstream of
+    the fused conditioning vector, i.e. the actual generative trunk of a
+    stage. Implementations are expected to expose `in_dim`/`out_dim`
+    attributes (as `ExpertTrunk`/`RoutedTrunk` do) — `giant.model.summary`
+    (gitea #46) reads them to report trunk widths without needing to know the
+    body architecture."""
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+
+class RoutedTrunk(Trunk):
+    def __init__(
+        self,
+        router: Router,
+        trunk_type: str,
+        in_dim: int,
+        out_dim: int,
+        hidden_dim: int,
+        n_res_blocks: int,
+        cond_dim: int,
+        dropout: float = 0.0,
+        block_conditioning: str = "add",
+    ) -> None:
+        super().__init__()
+        self.router = router
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.experts = nn.ModuleList(
+            [
+                build_expert_body(
+                    trunk_type,
+                    in_dim,
+                    out_dim,
+                    hidden_dim,
+                    n_res_blocks,
+                    cond_dim,
+                    dropout,
+                    block_conditioning=block_conditioning,
+                )
+                for _ in range(router.n_experts)
+            ]
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+    ) -> torch.Tensor:
+        return _route_forward(self.experts, self.router, x, cond, cond_cont, cond_cat, self.training)
+
+
+def build_trunk(
+    router: Router | None,
+    trunk_type: str,
+    in_dim: int,
+    out_dim: int,
+    hidden_dim: int,
+    n_res_blocks: int,
+    cond_dim: int,
+    dropout: float = 0.0,
+    block_conditioning: str = "add",
+) -> nn.Module:
+    """Build a stage's trunk: `trunk_type` (a `TRUNK_REGISTRY` key, e.g.
+    `"resmlp"`) selects the expert body architecture; `router`, if given,
+    wraps `router.n_experts` instances of that body in a `RoutedTrunk`
+    mixture — otherwise a single body is returned directly (no wrapper
+    class), which is what makes an unrouted trunk's state-dict keys land
+    directly under `trunk.*` instead of `trunk.experts.0.*` (see
+    `giant.model._legacy.migrate_legacy_state_dict`, which assumes exactly
+    this flat layout for a v0.2 monolithic checkpoint). `block_conditioning`
+    (a `BLOCK_REGISTRY` key, e.g. `"add"`/`"film"`/`"adaln"`) selects each
+    body's conditioning-injection mechanism (gitea #34).
+    """
+    if router is not None:
+        return RoutedTrunk(
+            router, trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout, block_conditioning
+        )
+    return build_expert_body(
+        trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout, block_conditioning
+    )
@@ -22,21 +22,31 @@ def gradient_penalty(
    norm to 1 — `x_hat`/`grad` are forced to all-zero for such a row, which
    would otherwise contribute a constant `(||0|| - 1)^2 == 1` bias to the
    mean regardless of critic behavior — so they're excluded from the mean.
+
+    Deliberately kept fp32 (`torch.autocast(..., enabled=False)`) regardless
+    of the caller's ambient `train.precision` autocast region: this is a
+    `create_graph=True` double-backward, and `grad.norm(2, dim=1)` sums
+    squares over the critic's full input width (hundreds of dims for stage
+    2), which overflows bf16's range at gradient magnitudes well within
+    normal early-WGAN-GP territory. Disclosed cost: the critic forward
+    inside this function always runs fp32, even when the rest of the WGAN
+    stage's step is bf16 (gitea #47).
    """
-    eps = torch.rand(real.size(0), 1, device=real.device)
-    x_hat = eps * real + (1 - eps) * fake
-    if mask is not None:
-        x_hat = x_hat * mask
-    x_hat = x_hat.requires_grad_(True)
-    scores = critic_fn(x_hat)
-    grad = torch.autograd.grad(outputs=scores.sum(), inputs=x_hat, create_graph=True)[0]
-    if mask is not None:
-        grad = grad * mask
-    penalty = (grad.norm(2, dim=1) - 1) ** 2
-    if mask is not None:
-        valid = (mask.sum(dim=1) > 0).float()
-        return (penalty * valid).sum() / valid.sum().clamp_min(1.0)
-    return penalty.mean()
+    with torch.autocast(real.device.type, enabled=False):
+        eps = torch.rand(real.size(0), 1, device=real.device)
+        x_hat = eps * real.float() + (1 - eps) * fake.float()
+        if mask is not None:
+            x_hat = x_hat * mask
+        x_hat = x_hat.requires_grad_(True)
+        scores = critic_fn(x_hat)
+        grad = torch.autograd.grad(outputs=scores.sum(), inputs=x_hat, create_graph=True)[0]
+        if mask is not None:
+            grad = grad * mask
+        penalty = (grad.norm(2, dim=1) - 1) ** 2
+        if mask is not None:
+            valid = (mask.sum(dim=1) > 0).float()
+            return (penalty * valid).sum() / valid.sum().clamp_min(1.0)
+        return penalty.mean()


 def critic_loss(
@@ -8,8 +8,8 @@ nuclide table doesn't cover — confirmed necessary for ~32% of the nuclear code
 actually present in the multi-material dataset
 (`0932fb02-f2ce-43ca-a4ef-60a2b1221bbc.parquet`).

-Also holds the v0.3.0 stage-2 categorical-type rollout decode (§3.3/§8/§11.3
-of docs/v0.3.0-design.md): `decode_topn_class`/`decode_embedding_nearest` turn
+Also holds the v0.3.0 stage-2 categorical-type rollout decode:
+`decode_topn_class`/`decode_embedding_nearest` turn
 `Stage2Autoregressive`/`Stage2OneShot`'s `"onehot"`/`"embedding"` type
 predictions back into concrete PDG codes, the one place a secondary's
 categorical/continuous type representation is ever discretized (its
@@ -58,9 +58,7 @@ def particle_mass_charge(pdg: int) -> tuple[float, float]:
        if _pdgid.is_nucleus(pdg):
            z, a = _pdgid.Z(pdg), _pdgid.A(pdg)
            if z is None or a is None:
-                raise ValueError(
-                    f"PDG {pdg}: is_nucleus but Z/A decode failed"
-                ) from None
+                raise ValueError(f"PDG {pdg}: is_nucleus but Z/A decode failed") from None
            return float(a) * _AMU_MEV, float(z)
        raise ValueError(
            f"PDG code {pdg} could not be resolved via the `particle` package "
@@ -109,9 +107,7 @@ def nearest_known_pdg(mass: np.ndarray, charge: np.ndarray, candidates) -> np.nd
    if len(resolved) == 0:
        raise ValueError("nearest_known_pdg: no resolvable candidates")
    codes = np.array([r[0] for r in resolved], dtype=np.int64)
-    table_log_mass = np.log(
-        np.array([r[1] for r in resolved], dtype=np.float64) + _LOG_EPS
-    )
+    table_log_mass = np.log(np.array([r[1] for r in resolved], dtype=np.float64) + _LOG_EPS)
    table_charge = np.array([r[2] for r in resolved], dtype=np.float64)

    mass = np.asarray(mass, dtype=np.float64)
@@ -144,9 +140,8 @@ def decode_topn_class(
    other_policy: str = "sample",
    rng: np.random.Generator | None = None,
 ) -> np.ndarray:
-    """`conditioning.particle.type` / `stage2_model.particle_type.target =
-    "onehot"` inference decode (docs/v0.3.0-design.md §3.3): per-row top-N
-    class index -> concrete PDG code.
+    """`stage2_model.particle_type.target = "onehot"` inference decode:
+    per-row top-N class index -> concrete secondary-species PDG code.

    class_idx: int array, any shape, values in `[0, n_classes)`.
    topn_map: the `TopNMap` (`giant.data.loader.build_pdg_topn_map_from_files`)
@@ -154,8 +149,11 @@ def decode_topn_class(
        except at the shared "other" index) plus `other_members` (the
        empirical within-"other" distribution, needed for `other_policy =
        "sample"`/`"modal"`).
-    n_classes: `conditioning.particle.emb_dim` — the class count; the "other"
-        bucket is index `n_classes - 1` by construction
+    n_classes: the resolved secondary-species class count
+        (`giant.model.models.resolve_type_n_classes` —
+        `stage2_model.particle_type.n_classes`, 0 = inherit
+        `conditioning.particle.emb_dim`; see gitea #29); the "other" bucket
+        is index `n_classes - 1` by construction
        (`giant.data.loader._topn_plus_other_map`).
    other_policy: `"sample"` draws from `other_members`' empirical frequency;
        `"modal"` always the single most common "other" member; `"drop"`
@@ -180,10 +178,7 @@ def decode_topn_class(
    n_other = int(other_mask.sum())
    if n_other:
        if not topn_map.other_members:
-            raise ValueError(
-                "decode_topn_class: 'other' class predicted but "
-                "topn_map.other_members is empty"
-            )
+            raise ValueError("decode_topn_class: 'other' class predicted but topn_map.other_members is empty")
        members = np.array(list(topn_map.other_members.keys()), dtype=np.int64)
        counts = np.array(list(topn_map.other_members.values()), dtype=np.float64)
        if other_policy == "drop":
@@ -205,9 +200,9 @@ def decode_embedding_nearest(
    emb_weight: np.ndarray,
    idx_to_pdg: dict[int, int],
 ) -> tuple[np.ndarray, np.ndarray]:
-    """`stage2_model.particle_type.target = "embedding"` inference decode
-    (docs/v0.3.0-design.md §3.3): L1-nearest row of the conditioning's own
-    particle embedding table, since a generative model's continuous output
+    """`stage2_model.particle_type.target = "embedding"` inference decode:
+    L1-nearest row of the conditioning's own particle embedding table, since
+    a generative model's continuous output
    essentially never lands within float tolerance of a table row (the exact-
    match form is only valid as a round-trip test assertion, never here).

@@ -220,9 +215,9 @@ def decode_embedding_nearest(
    idx_to_pdg: `invert_dense_map(pdg_map)` — embedding row index -> PDG.

    Returns `(pdg, l1_dist)`, both shaped like `vectors.shape[:-1]`. `l1_dist`
-    is the §11.3 diagnostic: a heavy tail means the decoder is emitting
-    vectors off the embedding manifold, the direct analogue of the species-
-    collapse symptom this redesign exists to fix.
+    is a diagnostic: a heavy tail means the decoder is emitting vectors off
+    the embedding manifold, the direct analogue of the species-collapse
+    symptom this redesign exists to fix.
    """
    emb_dim = vectors.shape[-1]
    flat = np.asarray(vectors, dtype=np.float64).reshape(-1, emb_dim)
@@ -31,8 +31,8 @@ from giant.data.transforms import (
    sorted_membership,
 )
 from giant.data.dataset import make_event_split, StreamingStepsDataset
-from giant.model.network import build_models, build_critics
-from giant.train import train as run_training
+from giant.model.network import build_models, build_critics, resolve_type_n_classes
+from giant.training import train as run_training


@dataclass
@@ -49,6 +49,7 @@ class SetupStageResult:
    mat_map: dict[str, int]
    proc_map: dict[str, int] | None
    pdg_topn_map: TopNMap | None
+    sec_type_topn_map: TopNMap | None
    mat_topn_map: TopNMap | None
    cond_norm: Normalizer
    tgt_norm: Normalizer
@@ -68,26 +69,19 @@ def _seed_energy_router(
    """Mutate `router_cfg["centers_init"]` in place from real data quantiles,
    when this stage's router is an enabled EnergyRouter. Shared by both
    stages' router configs — each seeded independently, since v0.3.0 stages
-    may have entirely different router configs (see docs/v0.3.0-design.md)."""
+    may have entirely different router configs."""
    active = router_cfg.get("enabled") and router_cfg.get("type") == "energy"
    if not active:
        return
    if energy_quantiles.size == 0:
-        echo(
-            "  warning: no energy samples collected — EnergyRouter falls back to "
-            "default centers"
-        )
+        echo("  warning: no energy samples collected — EnergyRouter falls back to default centers")
        return
    assert cond_norm.mean is not None and cond_norm.std is not None
    levels = np.linspace(0.0, 1.0, router_cfg["n_experts"])
    raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels)
-    centers_init = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[
-        energy_idx
-    ]
+    centers_init = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[energy_idx]
    router_cfg["centers_init"] = centers_init.astype(np.float32).tolist()
-    echo(
-        f"  seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}"
-    )
+    echo(f"  seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}")


 def run_setup_stage(
@@ -107,7 +101,7 @@ def run_setup_stage(
    `cfg` is the full merged v0.3 config (`conditioning`/`stage1_model`/
    `stage2_model`), already passed through `giant.config.validate_config`.
    `conditioning.particle.type` and `conditioning.material.type` are
-    independent (docs/v0.3.0-design.md §3.1) and may differ.
+    independent and may differ.

    Reads from and writes to the `giant.data.setup_cache` sidecar when
    `cache_setup` is set (`rebuild_setup_cache` ignores — but still
@@ -142,23 +136,14 @@ def run_setup_stage(
        if cache is not None:
            cache.event_index = (unique_ids, counts)

-    train_events, val_events = make_event_split(
-        unique_ids, val_fraction=val_fraction, seed=seed
-    )
+    train_events, val_events = make_event_split(unique_ids, val_fraction=val_fraction, seed=seed)
    events_arr = np.array(sorted(train_events))
    n_train_steps = setup_cache.n_train_steps_for_split(unique_ids, counts, events_arr)
-    echo(
-        f"  {int(counts.sum()):,} steps  |  "
-        f"{len(train_events)} train events  |  "
-        f"{len(val_events)} val events"
-    )
+    echo(f"  {int(counts.sum()):,} steps  |  {len(train_events)} train events  |  {len(val_events)} val events")

    if cache is not None and cache.vocab is not None:
        pdg_map, mat_map = cache.vocab
-        echo(
-            f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, "
-            f"{len(mat_map)} materials)"
-        )
+        echo(f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, {len(mat_map)} materials)")
    else:
        echo("building vocabulary maps …")
        pdg_map, mat_map = build_index_maps_from_files(files)
@@ -174,11 +159,7 @@ def run_setup_stage(
    # need that generality.
    proc_map: dict[str, int] | None = None
    process_router_cfg = next(
-        (
-            r
-            for r in (stage1_router, stage2_router)
-            if r.get("enabled") and r.get("type") == "process"
-        ),
+        (r for r in (stage1_router, stage2_router) if r.get("enabled") and r.get("type") == "process"),
        None,
    )
    if process_router_cfg is not None:
@@ -186,10 +167,7 @@ def run_setup_stage(
        cached_proc_map = cache.proc_maps.get(n_experts) if cache is not None else None
        if cached_proc_map is not None:
            proc_map = cached_proc_map
-            echo(
-                f"process vocabulary: cache hit ({len(proc_map)} labels, "
-                f"{n_experts} experts)"
-            )
+            echo(f"process vocabulary: cache hit ({len(proc_map)} labels, {n_experts} experts)")
        else:
            echo("building process vocabulary …")
            proc_map = build_process_map_from_files(files, n_experts=n_experts)
@@ -197,35 +175,43 @@ def run_setup_stage(
            if cache is not None:
                cache.proc_maps[n_experts] = proc_map

-    # Top-N-plus-other maps for onehot conditioning/type axes
-    # (docs/v0.3.0-design.md §8). The PDG axis is shared by
-    # conditioning.particle.type="onehot" and
-    # stage2_model.particle_type.target="onehot" (both key off
-    # conditioning.particle.emb_dim), so at most one PDG scan is needed even
-    # if both consumers are active. The material axis is independent.
+    # Top-N-plus-other maps for onehot conditioning/type axes.
+    # The PDG axis is used independently by conditioning.particle.type="onehot"
+    # (cond_cat's onehot feature) and stage2_model.particle_type.target="onehot"
+    # (secondary-species decode) — their class counts can now differ (gitea
+    # #29: stage2_model.particle_type.n_classes, 0 = inherit
+    # conditioning.particle.emb_dim), so each is resolved and built
+    # independently via _pdg_topn below. cache.topn_maps is keyed by
+    # (axis, n_classes) (setup_cache.topn_key), so when the two resolve to
+    # the same N the second call is a cache hit against the first — no extra
+    # scan in the common case where they still match. The material axis is
+    # independent of both.
    particle_cfg = cfg["conditioning"]["particle"]
    material_cfg = cfg["conditioning"]["material"]
-    particle_type_target = cfg["stage2_model"].get("particle_type", {}).get("target")
+    particle_type_cfg = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type"))
+    particle_type_target = particle_type_cfg.target

-    pdg_topn_map: TopNMap | None = None
-    if particle_cfg["type"] == "onehot" or particle_type_target == "onehot":
-        n_classes = particle_cfg["emb_dim"]
+    def _pdg_topn(n_classes: int) -> TopNMap:
        cache_key = setup_cache.topn_key("pdg", n_classes)
        cached = cache.topn_maps.get(cache_key) if cache is not None else None
        if cached is not None:
-            pdg_topn_map = cached
-            echo(
-                f"pdg top-N map: cache hit ({len(pdg_topn_map.class_map)} codes, "
-                f"{n_classes} classes)"
-            )
-        else:
-            echo("building pdg top-N map …")
-            pdg_topn_map = build_pdg_topn_map_from_files(files, n_classes=n_classes)
-            echo(
-                f"  {len(pdg_topn_map.class_map)} pdg codes mapped to {n_classes} classes"
-            )
-            if cache is not None:
-                cache.topn_maps[cache_key] = pdg_topn_map
+            echo(f"pdg top-N map: cache hit ({len(cached.class_map)} codes, {n_classes} classes)")
+            return cached
+        echo("building pdg top-N map …")
+        topn_map = build_pdg_topn_map_from_files(files, n_classes=n_classes)
+        echo(f"  {len(topn_map.class_map)} pdg codes mapped to {n_classes} classes")
+        if cache is not None:
+            cache.topn_maps[cache_key] = topn_map
+        return topn_map
+
+    pdg_topn_map: TopNMap | None = None
+    if particle_cfg["type"] == "onehot":
+        pdg_topn_map = _pdg_topn(particle_cfg["emb_dim"])
+
+    sec_type_topn_map: TopNMap | None = None
+    if particle_type_target == "onehot":
+        sec_type_n_classes = resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
+        sec_type_topn_map = _pdg_topn(sec_type_n_classes)

    mat_topn_map: TopNMap | None = None
    if material_cfg["type"] == "onehot":
@@ -234,29 +220,17 @@ def run_setup_stage(
        cached = cache.topn_maps.get(cache_key) if cache is not None else None
        if cached is not None:
            mat_topn_map = cached
-            echo(
-                f"material top-N map: cache hit ({len(mat_topn_map.class_map)} "
-                f"materials, {n_classes} classes)"
-            )
+            echo(f"material top-N map: cache hit ({len(mat_topn_map.class_map)} materials, {n_classes} classes)")
        else:
            echo("building material top-N map …")
-            mat_topn_map = build_topn_map_from_files(
-                files, "material", n_classes=n_classes, cast=str
-            )
-            echo(
-                f"  {len(mat_topn_map.class_map)} materials mapped to {n_classes} classes"
-            )
+            mat_topn_map = build_topn_map_from_files(files, "material", n_classes=n_classes, cast=str)
+            echo(f"  {len(mat_topn_map.class_map)} materials mapped to {n_classes} classes")
            if cache is not None:
                cache.topn_maps[cache_key] = mat_topn_map

-    energy_router_active = any(
-        r.get("enabled") and r.get("type") == "energy"
-        for r in (stage1_router, stage2_router)
-    )
+    energy_router_active = any(r.get("enabled") and r.get("type") == "energy" for r in (stage1_router, stage2_router))
    energy_idx = 3
-    norm_key = setup_cache.normalizer_key(
-        val_fraction, seed, particle_conditioning, material_conditioning
-    )
+    norm_key = setup_cache.normalizer_key(val_fraction, seed, particle_conditioning, material_conditioning)
    entry = cache.normalizers.get(norm_key) if cache is not None else None

    if entry is not None:
@@ -281,28 +255,35 @@ def run_setup_stage(
        # router against this same (val_fraction, seed, conditioning) key
        # never needs to rescan just to seed centers.
        collect_energy_sample = energy_router_active or cache is not None
-        energy_sampler = (
-            _ReservoirSampler(capacity=100_000) if collect_energy_sample else None
-        )
+        energy_sampler = _ReservoirSampler(capacity=100_000) if collect_energy_sample else None
        for i, path in enumerate(files):
            for chunk in iter_file_chunks(path, offset=event_id_offset(i), k_max=k_max):
                mask = sorted_membership(chunk["event_id"], events_arr)
                if not mask.any():
                    continue
                chunk_tr = {k: v[mask] for k, v in chunk.items()}
-                cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _, _ = (
-                    build_features(
-                        chunk_tr,
-                        pdg_map,
-                        mat_map,
-                        proc_map=proc_map,
-                        require_secondaries=True,
-                        particle_conditioning=particle_conditioning,
-                        material_conditioning=material_conditioning,
-                        sec_phys_only=True,
-                        k_max=k_max,
-                    )
+                feats = build_features(
+                    chunk_tr,
+                    pdg_map,
+                    mat_map,
+                    proc_map=proc_map,
+                    require_secondaries=True,
+                    particle_conditioning=particle_conditioning,
+                    material_conditioning=material_conditioning,
+                    sec_phys_only=True,
+                    # This pass reads only cond_cont/sec_cont, never cond_cat —
+                    # but cond_cat's width is the conditioning modes' call
+                    # (giant.cond_layout.CondLayout), so an "onehot" axis still
+                    # has to be handed its map rather than silently yielding a
+                    # narrower array.
+                    pdg_topn_map=pdg_topn_map.class_map if pdg_topn_map is not None else None,
+                    mat_topn_map=mat_topn_map.class_map if mat_topn_map is not None else None,
+                    k_max=k_max,
                )
+                cond_cont = feats.cond_cont
+                target_s1 = feats.target_s1
+                n_sec = feats.n_sec
+                sec_cont = feats.sec_cont
                cond_acc.update(cond_cont)
                tgt_acc.update(target_s1)
                if energy_sampler is not None:
@@ -336,6 +317,7 @@ def run_setup_stage(
        mat_map=mat_map,
        proc_map=proc_map,
        pdg_topn_map=pdg_topn_map,
+        sec_type_topn_map=sec_type_topn_map,
        mat_topn_map=mat_topn_map,
        cond_norm=cond_norm,
        tgt_norm=tgt_norm,
@@ -377,7 +359,7 @@ def run_train_job(
            "section)"
        )

-    config.validate_config(cfg)
+    config.validate_config(cfg, resume=resume is not None)
    particle_conditioning = cfg["conditioning"]["particle"]["type"]
    material_conditioning = cfg["conditioning"]["material"]["type"]
    k_max = cfg["stage2_model"]["k_max"]
@@ -403,11 +385,11 @@ def run_train_job(
        setup.n_train_steps,
    )

-    # cond_cat's onehot columns (docs/v0.3.0-design.md decision 4) are
-    # present per-axis, independently, under that axis's own
-    # conditioning.{particle,material}.type == "onehot" (§3.1: the two axes
-    # may mix freely). run_setup_stage builds each map whenever its own axis
-    # is "onehot" (see its own particle_cfg["type"]/material_cfg["type"]
+    # cond_cat's onehot columns are present per-axis, independently, under
+    # that axis's own conditioning.{particle,material}.type == "onehot"
+    # (the two axes may mix freely). run_setup_stage builds each map
+    # whenever its own axis is "onehot" (see its own
+    # particle_cfg["type"]/material_cfg["type"]
    # checks), so they're guaranteed non-None here — asserted, not just
    # assumed, so a future wiring bug fails loudly instead of silently
    # dropping the onehot columns.
@@ -422,13 +404,11 @@ def run_train_job(

    # The secondary type-index map depends on stage2_model.particle_type.target,
    # independently of conditioning's own onehot/embedding choice above
-    # (docs/v0.3.0-design.md §3.3 — physical stays untouched/None).
-    particle_type_target = (
-        cfg["stage2_model"].get("particle_type", {}).get("target", "physical")
-    )
+    # (physical stays untouched/None).
+    particle_type_target = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type")).target
    if particle_type_target == "onehot":
-        assert setup.pdg_topn_map is not None
-        sec_type_class_map = setup.pdg_topn_map.class_map
+        assert setup.sec_type_topn_map is not None
+        sec_type_class_map = setup.sec_type_topn_map.class_map
    elif particle_type_target == "embedding":
        sec_type_class_map = pdg_map
    else:
@@ -455,6 +435,7 @@ def run_train_job(
        mat_topn_map=cond_mat_topn,
        sec_type_class_map=sec_type_class_map,
        k_max=k_max,
+        seed=t["seed"],
    )
    val_ds = StreamingStepsDataset(
        files=files,
@@ -532,6 +513,7 @@ def run_train_job(
        mat_map={str(k): v for k, v in mat_map.items()},
        proc_map=proc_map,
        pdg_topn_map=setup.pdg_topn_map,
+        sec_type_topn_map=setup.sec_type_topn_map,
        mat_topn_map=setup.mat_topn_map,
        model_config=model_config,
        resume_path=resume,
@@ -53,14 +53,14 @@ if TYPE_CHECKING:


 class L1DistCollector:
-    """Accumulates the §11.3 L1-distance diagnostic across a whole rollout
+    """Accumulates the L1-distance diagnostic across a whole rollout
    run: the L1 distance between each emitted secondary's raw predicted
    embedding vector and the nearest table row it snapped to (only
    meaningful under `particle_type.target = "embedding"` —
    `giant.particles.decode_embedding_nearest`). A heavy tail means the
    decoder is emitting vectors off the embedding manifold — the direct
    analogue of the species-collapse symptom the v0.3.0 redesign exists to
-    fix (docs/v0.3.0-design.md §11.3).
+    fix.

    Not folded into `rollout()`'s own return value (which is shape-typed as
    step records, see `_RECORD_KEYS`/`RolloutSummary`) — passed in and read
@@ -117,15 +117,12 @@ def decode_secondary_identity(
    pre_dir: np.ndarray,
    sec_phys_norm: Normalizer,
    pdg_map: dict[int, int],
-    pdg_topn_map: "TopNMap | None",
+    sec_type_topn_map: "TopNMap | None",
    other_policy: str,
    rng: np.random.Generator | None,
-) -> tuple[
-    np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray | None
-]:
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray | None]:
    """Decode Stage 2's raw (sec_cont, sec_type) output into physical
-    secondary attributes, branching on `sec_decoder.particle_type_cfg`
-    (docs/v0.3.0-design.md §3.3):
+    secondary attributes, branching on `sec_decoder.particle_type_cfg`:

    - `"physical"`: unchanged v0.2 path — `sec_type` already *is* (log_mass,
      charge), used as the secondary's identity as-is (no snapping).
@@ -138,40 +135,38 @@ def decode_secondary_identity(
    - `"embedding"`: `sec_type` is a raw vector in the conditioning's own
      embedding space — `giant.particles.decode_embedding_nearest` L1-snaps
      it to the nearest table row for the PDG (+ physics via
-      `particle_phys_array`), and also returns the L1 distance (§11.3
-      diagnostic — see `giant/rollout.py`'s L1-distance accumulator).
+      `particle_phys_array`), and also returns the L1 distance (see this
+      module's `L1DistCollector`).

    Returns (sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg,
    sec_type_l1_dist) — the last is `None` except under `"embedding"`.
    """
-    target = sec_decoder.particle_type_cfg.get("target", "physical")
+    target = sec_decoder.particle_type_cfg.target

    if target == "physical":
        sec_full = torch.cat([sec_cont, sec_type], dim=-1).cpu().numpy()
        sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid = decode_secondaries(
            sec_full, n_sec_np, e_sec, pre_dir, sec_phys_normalizer=sec_phys_norm
        )
-        sec_pdg = nearest_known_pdg(
-            sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()
-        ).reshape(sec_mass.shape)
+        sec_pdg = nearest_known_pdg(sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()).reshape(
+            sec_mass.shape
+        )
        return sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg, None

-    sec_E, sec_dir_world, sec_valid = decode_secondary_cont(
-        sec_cont.cpu().numpy(), n_sec_np, e_sec, pre_dir
-    )
+    sec_E, sec_dir_world, sec_valid = decode_secondary_cont(sec_cont.cpu().numpy(), n_sec_np, e_sec, pre_dir)
    sec_type_np = sec_type.cpu().numpy()
    l1_dist = None

    if target == "onehot":
-        if pdg_topn_map is None:
+        if sec_type_topn_map is None:
            raise RuntimeError(
-                "particle_type.target='onehot' rollout needs pdg_topn_map "
-                "(the checkpoint's saved top-N map) — see ckpt['pdg_topn_map']"
+                "particle_type.target='onehot' rollout needs sec_type_topn_map "
+                "(the checkpoint's saved top-N map) — see ckpt['sec_type_topn_map']"
            )
        class_idx = sec_type_np.argmax(axis=-1)
        sec_pdg = decode_topn_class(
            class_idx,
-            pdg_topn_map,
+            sec_type_topn_map,
            n_classes=sec_decoder.type_dim,
            other_policy=other_policy,
            rng=rng,
@@ -183,12 +178,8 @@ def decode_secondary_identity(
        l1_dist = np.where(sec_valid, l1_dist, 0.0).astype(np.float32)

    sec_mass, sec_charge = particle_phys_array(sec_pdg.reshape(-1)).T
-    sec_mass = np.where(sec_valid, sec_mass.reshape(sec_pdg.shape), 0.0).astype(
-        np.float32
-    )
-    sec_charge = np.where(sec_valid, sec_charge.reshape(sec_pdg.shape), 0.0).astype(
-        np.float32
-    )
+    sec_mass = np.where(sec_valid, sec_mass.reshape(sec_pdg.shape), 0.0).astype(np.float32)
+    sec_charge = np.where(sec_valid, sec_charge.reshape(sec_pdg.shape), 0.0).astype(np.float32)
    sec_pdg = np.where(sec_valid, sec_pdg, 0).astype(np.int64)
    return sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg, l1_dist

@@ -300,13 +291,9 @@ class _Recorder:
    RAM until the very end.
    """

-    def __init__(
-        self, sink: Callable[[dict[str, np.ndarray]], None] | None = None
-    ) -> None:
+    def __init__(self, sink: Callable[[dict[str, np.ndarray]], None] | None = None) -> None:
        self._sink = sink
-        self._cols: dict[str, list] | None = (
-            None if sink is not None else {k: [] for k in _RECORD_KEYS}
-        )
+        self._cols: dict[str, list] | None = None if sink is not None else {k: [] for k in _RECORD_KEYS}
        self.n_rows = 0
        self.termination_reason_counts: Counter[str] = Counter()

@@ -314,10 +301,7 @@ class _Recorder:
        n = len(cols["event_id"])
        if n == 0:
            return
-        row = {
-            k: np.asarray(cols[k], dtype=_RECORD_DTYPES[k]).reshape(n)
-            for k in _RECORD_KEYS
-        }
+        row = {k: np.asarray(cols[k], dtype=_RECORD_DTYPES[k]).reshape(n) for k in _RECORD_KEYS}
        self.n_rows += n
        reasons = row["termination_reason"]
        nonempty = reasons[reasons != ""]
@@ -334,8 +318,7 @@ class _Recorder:

    def to_dict(self) -> dict[str, np.ndarray]:
        assert self._cols is not None, (
-            "to_dict() is unavailable when streaming to a sink — use "
-            "n_rows/termination_reason_counts instead"
+            "to_dict() is unavailable when streaming to a sink — use n_rows/termination_reason_counts instead"
        )
        out = {}
        for k, chunks in self._cols.items():
@@ -461,6 +444,7 @@ def rollout(
    material_conditioning: str = "embedding",
    pdg_topn_map: "TopNMap | None" = None,
    mat_topn_map: "TopNMap | None" = None,
+    sec_type_topn_map: "TopNMap | None" = None,
    other_policy: str = "sample",
    seed: int | None = None,
    stage1_ddpm_steps: int = 1000,
@@ -480,24 +464,25 @@ def rollout(
    `n_events * max_steps * avg_tracks_per_event`.

    There is no `mode` parameter — each stage's generative objective is read
-    directly off the model instance's own `generator_kind`
-    (docs/v0.3.0-design.md decision 2: stage 1 and stage 2 objectives are
-    independent, e.g. `stage1_model.generator="flow"` +
-    `stage2_model.generator="wgan"`), and the decoder (one-shot vs
+    directly off the model instance's own `generator_kind` (stage 1 and
+    stage 2 objectives are independent, e.g. `stage1_model.generator="flow"`
+    + `stage2_model.generator="wgan"`), and the decoder (one-shot vs
    autoregressive) is inferred from `sec_decoder`'s own class — see
    `sample_stage1`/`sample_stage2` (giant.sample).

-    `pdg_topn_map`/`mat_topn_map` serve two independent purposes that happen
-    to share `pdg_topn_map` (docs/v0.3.0-design.md §8 — one PDG map, not
-    two): they're required whenever `particle_conditioning`/
-    `material_conditioning` is `"onehot"` (feeds `build_cond_features`'s
-    extra `cond_cat` top-N columns — §3.1), and `pdg_topn_map`/`other_policy`
-    are additionally read under `stage2_model.particle_type.target =
-    "onehot"` (§3.3, secondary-species decode). `seed` seeds the
-    `other_policy = "sample"` draw only (torch/numpy sampling itself is
-    seeded by the caller, same as today).
+    `pdg_topn_map`/`mat_topn_map`/`sec_type_topn_map` serve three independent
+    purposes, no longer required to share one map (see gitea #29):
+    `pdg_topn_map`/`mat_topn_map` are required whenever
+    `particle_conditioning`/`material_conditioning` is `"onehot"` (feeds
+    `build_cond_features`'s extra `cond_cat` top-N columns); `sec_type_topn_map`/
+    `other_policy` are required instead under
+    `stage2_model.particle_type.target = "onehot"` (secondary-species
+    decode) — its class count (`stage2_model.particle_type.n_classes`) may
+    differ from `pdg_topn_map`'s. `seed` seeds the `other_policy = "sample"`
+    draw only (torch/numpy sampling itself is seeded by the caller, same as
+    today).

-    `l1_dist_collector`, if given, accumulates the §11.3 embedding-distance
+    `l1_dist_collector`, if given, accumulates the embedding-distance
    diagnostic across the whole run — see `L1DistCollector`. Only populated
    under `particle_type.target = "embedding"`; a no-op otherwise.
    """
@@ -506,6 +491,11 @@ def rollout(
            "conditioning.particle.type='onehot' rollout needs pdg_topn_map "
            "(the checkpoint's saved top-N map) — see ckpt['pdg_topn_map']"
        )
+    if sec_decoder.particle_type_cfg.target == "onehot" and sec_type_topn_map is None:
+        raise RuntimeError(
+            "stage2_model.particle_type.target='onehot' rollout needs sec_type_topn_map "
+            "(the checkpoint's saved top-N map) — see ckpt['sec_type_topn_map']"
+        )
    if material_conditioning == "onehot" and mat_topn_map is None:
        raise RuntimeError(
            "conditioning.material.type='onehot' rollout needs mat_topn_map "
@@ -555,6 +545,7 @@ def rollout(
                    material_conditioning,
                    pdg_topn_map,
                    mat_topn_map,
+                    sec_type_topn_map,
                    other_policy,
                    rng,
                    stage1_ddpm_steps,
@@ -593,6 +584,7 @@ def _step_chunk(
    material_conditioning,
    pdg_topn_map,
    mat_topn_map,
+    sec_type_topn_map,
    other_policy,
    rng,
    stage1_ddpm_steps,
@@ -622,33 +614,19 @@ def _step_chunk(
    # --- Pre-step termination gates (in priority order; each track picks one) ---
    stop = np.zeros(n, dtype=bool)
    escaped_sel = escaped & ~stop
-    rec.add(
-        **_terminal_rows(
-            tr, escaped_sel, TERM_ESCAPED, edep=np.zeros(int(escaped_sel.sum()))
-        )
-    )
+    rec.add(**_terminal_rows(tr, escaped_sel, TERM_ESCAPED, edep=np.zeros(int(escaped_sel.sum()))))
    stop |= escaped_sel

    unknown_sel = ~known_pdg & ~stop
-    rec.add(
-        **_terminal_rows(
-            tr, unknown_sel, TERM_UNKNOWN_PDG, edep=tr["pre_E"][unknown_sel]
-        )
-    )
+    rec.add(**_terminal_rows(tr, unknown_sel, TERM_UNKNOWN_PDG, edep=tr["pre_E"][unknown_sel]))
    stop |= unknown_sel

    cutoff_sel = (tr["pre_E"] < energy_cutoff) & ~stop
-    rec.add(
-        **_terminal_rows(
-            tr, cutoff_sel, TERM_ENERGY_CUTOFF, edep=tr["pre_E"][cutoff_sel]
-        )
-    )
+    rec.add(**_terminal_rows(tr, cutoff_sel, TERM_ENERGY_CUTOFF, edep=tr["pre_E"][cutoff_sel]))
    stop |= cutoff_sel

    maxstep_sel = (tr["step_in_track"] >= max_steps) & ~stop
-    rec.add(
-        **_terminal_rows(tr, maxstep_sel, TERM_MAX_STEPS, edep=tr["pre_E"][maxstep_sel])
-    )
+    rec.add(**_terminal_rows(tr, maxstep_sel, TERM_MAX_STEPS, edep=tr["pre_E"][maxstep_sel]))
    stop |= maxstep_sel

    active = ~stop
@@ -682,42 +660,27 @@ def _step_chunk(
        cond_norm,
        particle_conditioning=particle_conditioning,
        material_conditioning=material_conditioning,
-        pdg_topn_map=pdg_topn_map.class_map
-        if particle_conditioning == "onehot"
-        else None,
-        mat_topn_map=mat_topn_map.class_map
-        if material_conditioning == "onehot"
-        else None,
+        pdg_topn_map=pdg_topn_map.class_map if particle_conditioning == "onehot" else None,
+        mat_topn_map=mat_topn_map.class_map if material_conditioning == "onehot" else None,
    )
    cc = torch.from_numpy(cond_cont).float().to(device)
    ck = torch.from_numpy(cond_cat).long().to(device)

-    stage1_norm, n_sec_pred_stage1 = sample_stage1(
-        stage1_model, cc, ck, steps, stage1_ddpm_steps
-    )
+    stage1_norm, n_sec_pred_stage1 = sample_stage1(stage1_model, cc, ck, steps, stage1_ddpm_steps)
    raw = tgt_norm.inverse_transform(stage1_norm.cpu().numpy())

    step_length = inv_log_transform(raw[:, 0])
    edep, e_sec, post_E, _delta = energy_simplex_decode(raw[:, 1:3], tr["pre_E"])

    post_dir_local = raw[:, 3:6].copy()
-    post_dir_local /= np.clip(
-        np.linalg.norm(post_dir_local, axis=1, keepdims=True), 1e-8, None
-    )
+    post_dir_local /= np.clip(np.linalg.norm(post_dir_local, axis=1, keepdims=True), 1e-8, None)
    post_dir_world = inv_local_frame_rotation(tr["pre_dir"], post_dir_local)

    travel_dir_local = raw[:, 6:9].copy()
-    travel_dir_local /= np.clip(
-        np.linalg.norm(travel_dir_local, axis=1, keepdims=True), 1e-8, None
-    )
-    post_pos = reconstruct_post_pos(
-        tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local
-    )
+    travel_dir_local /= np.clip(np.linalg.norm(travel_dir_local, axis=1, keepdims=True), 1e-8, None)
+    post_pos = reconstruct_post_pos(tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local)

-    n_sec_pred = resolve_n_sec(
-        stage1_model, sec_decoder, cc, ck, stage1_norm, n_sec_pred_stage1
-    )
-    n_sec_np = n_sec_pred.cpu().numpy().astype(np.int64)
+    n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cc, ck, stage1_norm, n_sec_pred_stage1)

    # --- Secondaries ---
    # No snapping for "physical"/history-facing state elsewhere in the
@@ -727,23 +690,25 @@ def _step_chunk(
    # decode_secondary_identity's docstring for how each
    # particle_type.target differs on whether PDG resolution is a real
    # identity decision or just a reporting label.
-    sec_cont, sec_type, _valid = sample_stage2(
-        sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps
-    )
-    sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, sec_type_l1_dist = (
-        decode_secondary_identity(
-            sec_decoder,
-            sec_cont,
-            sec_type,
-            n_sec_np,
-            e_sec,
-            tr["pre_dir"],
-            sec_phys_norm,
-            pdg_map,
-            pdg_topn_map,
-            other_policy,
-            rng,
-        )
+    sec_cont, sec_type, sec_valid = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
+    # A stop-token decoder resolves n_sec_pred=None above — the real count
+    # only exists once sample_stage2 has actually generated (or stopped
+    # generating) tokens, so read it back off sec_valid here. Under every
+    # other n_sec.mode sec_valid was built FROM n_sec_pred, so this is a
+    # no-op round trip in those cases.
+    n_sec_np = sec_valid.sum(dim=-1).cpu().numpy().astype(np.int64)
+    sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, sec_type_l1_dist = decode_secondary_identity(
+        sec_decoder,
+        sec_cont,
+        sec_type,
+        n_sec_np,
+        e_sec,
+        tr["pre_dir"],
+        sec_phys_norm,
+        pdg_map,
+        sec_type_topn_map,
+        other_policy,
+        rng,
    )
    sec_valid = np.arange(sec_E.shape[1])[None, :] < n_sec_np[:, None]
    if l1_dist_collector is not None and sec_type_l1_dist is not None:
@@ -2,7 +2,7 @@ import torch
 import torch.nn.functional as F

 from giant.constants import CONT_SLOT_DIM, X_DIM
-from giant.model.network import Stage2Autoregressive, stage2_trunk_sec_dim
+from giant.model.network import DdpmObjective, Stage2Autoregressive, build_objective, stage2_trunk_sec_dim
 from giant.model.schedule import CosineSchedule


@@ -10,10 +10,9 @@ def _predict_n_sec_if_owned(
    model: torch.nn.Module, cond_cont: torch.Tensor, cond_cat: torch.Tensor
 ) -> torch.Tensor | None:
    """Stage-1 `n_sec_head` is only present on a migrated v0.2 checkpoint
-    (docs/v0.3.0-design.md decision 1 moves it to stage 2 for fresh runs —
-    see `Stage1Model`'s docstring). `None` here means "ask stage 2 instead",
-    which every caller (`giant/rollout.py`, `giant/cli.py`) must do for a
-    fresh checkpoint."""
+    (fresh runs move it to stage 2 — see `Stage1Model`'s docstring). `None`
+    here means "ask stage 2 instead", which every caller (`giant/rollout.py`,
+    `giant/cli.py`) must do for a fresh checkpoint."""
    if getattr(model, "n_sec_head", None) is None:
        return None
    logits = model.predict_n_sec(cond_cont, cond_cat)
@@ -68,9 +67,7 @@ def sample_ddpm(
        alpha = schedule.alphas[i]
        alpha_bar = schedule.alpha_bars[i]
        z = torch.randn_like(x) if i > 0 else torch.zeros_like(x)
-        x = (1.0 / alpha.sqrt()) * (
-            x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred
-        ) + beta.sqrt() * z
+        x = (1.0 / alpha.sqrt()) * (x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred) + beta.sqrt() * z
    return x, _predict_n_sec_if_owned(model, cond_cont, cond_cat)


@@ -124,7 +121,7 @@ def _stage2_flat_width(sec_decoder: torch.nn.Module) -> int:
    (continuous + type) under `particle_type.target = "physical"` or
    `generator = "wgan"`, continuous-only otherwise (the type slice then
    comes from `predict_type` instead — see `stage2_trunk_sec_dim`'s
-    docstring, docs/v0.3.0-design.md decision 2)."""
+    docstring)."""
    return stage2_trunk_sec_dim(
        sec_decoder.particle_type_cfg,
        sec_decoder.generator_kind,
@@ -134,8 +131,8 @@ def _stage2_flat_width(sec_decoder: torch.nn.Module) -> int:


 def _type_folded(sec_decoder: torch.nn.Module) -> bool:
-    target = sec_decoder.particle_type_cfg.get("target", "physical")
-    return target == "physical" or sec_decoder.generator_kind == "wgan"
+    target = sec_decoder.particle_type_cfg.target
+    return target == "physical" or build_objective(sec_decoder.generator_kind).folds_type_slice


 def _decode_stage2_flat(
@@ -147,9 +144,9 @@ def _decode_stage2_flat(
    n_sec_pred: torch.Tensor,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Reshape a flat `(B, flat_width)` `Stage2OneShot` output into per-slot
-    tensors, generator/`particle_type.target`-agnostic (docs/v0.3.0-design.md
-    decision 2/3): shared by `sample_secondaries`/`sample_secondaries_wgan`,
-    which differ only in how `x` was produced.
+    tensors, generator/`particle_type.target`-agnostic: shared by
+    `sample_secondaries`/`sample_secondaries_wgan`, which differ only in how
+    `x` was produced.

    Returns (sec_cont, sec_type, sec_valid):
        sec_cont:  (B, k_max, CONT_SLOT_DIM) — [stick_logit, local_dir]
@@ -174,9 +171,7 @@ def _decode_stage2_flat(
    else:
        sec_cont = x.view(B, k_max, CONT_SLOT_DIM)
        sec_type = sec_decoder.predict_type(cond_cont, cond_cat, stage1_out)
-    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(
-        1
-    )
+    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1)
    return sec_cont, sec_type, sec_valid


@@ -207,9 +202,7 @@ def sample_secondaries(
        v = sec_decoder(x, cond_cont, cond_cat, stage1_out, t=t)
        x = x + v * dt

-    return _decode_stage2_flat(
-        sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred
-    )
+    return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred)


@torch.no_grad()
@@ -227,9 +220,7 @@ def sample_secondaries_wgan(
    B = cond_cont.size(0)
    z = torch.randn(B, sec_decoder.noise_dim, device=cond_cont.device)
    x = sec_decoder(z, cond_cont, cond_cat, stage1_out)
-    return _decode_stage2_flat(
-        sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred
-    )
+    return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred)


@torch.no_grad()
@@ -238,26 +229,48 @@ def sample_secondaries_ar(
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    stage1_out: torch.Tensor,
-    n_sec_pred: torch.Tensor,
+    n_sec_pred: torch.Tensor | None,
    steps: int = 10,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """`Stage2Autoregressive` inference loop (docs/v0.3.0-design.md §6.4):
-    one token at a time, in descending-energy slot order, `k_max` sequential
-    calls. Unlike training (teacher forcing, §6.2 point 3 — a single
-    parallel pass over ground-truth tokens, see
-    `giant.train._assemble_stage2_ar_inputs`), there is no ground truth at
-    inference: each token's conditioning is built free-running, from the
-    PREVIOUS TOKEN'S OWN just-generated output — the train/inference gap
-    §6.2 point 4 explicitly flags as the cost of markov history's
+    """`Stage2Autoregressive` inference loop: one token at a time, in
+    descending-energy slot order, up to `k_max` sequential calls. Unlike
+    training (teacher forcing — a single parallel pass over ground-truth
+    tokens, see `giant.training.stage2_inputs._assemble_stage2_ar_inputs`),
+    there is no ground truth at inference: each token's conditioning is built
+    free-running, from the PREVIOUS TOKEN'S OWN just-generated output — the
+    train/inference gap that is the cost of markov history's
    expressiveness.

    A `{flow,ddpm}` token costs `steps` ODE substeps; `wgan` costs one pass —
-    §6.4's "K sequential forwards" cost note applies per-token here, not
+    the "K sequential forwards" cost applies per-token here, not
    once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per
-    physics step.
+    physics step (or ~`n_sec * steps` under `n_sec_pred=None` below, once
+    every row in the batch has stopped).
+
+    `n_sec_pred`, if given, fixes each row's secondary count up front (as
+    resolved by `resolve_n_sec` — `n_sec.mode` in `("head", "truth")`, or a
+    stop-token decoder driven by `_assemble_stage2_ar_inputs_scheduled`'s
+    ground-truth `n_sec`, which must run the *full* `k_max`-length free-
+    running self-sample regardless of the decoder's own stop head — the
+    scheduled-sampling training contract does not truncate). This always
+    runs the full `k_max`-iteration loop, masking by the given count at the
+    end exactly as before.
+
+    `n_sec_pred=None` is only valid when `sec_decoder.stop_head` is set
+    (`n_sec.mode = "stop_token"`): before generating each slot's token, that
+    slot's own stop logit (`predict_stop`, evaluated on the same prefix
+    conditioning as the token itself — see `predict_type`'s docstring for
+    why this needs no extra state) decides whether generation should have
+    already stopped, per `sec_decoder.stop_sampling` ("greedy": threshold at
+    0; "sample": a Bernoulli draw at `sigmoid(logit)`). A row's own
+    `n_sec_pred` is the first slot index where this fires; once every row in
+    the batch has fired, the loop breaks before spending a model call on the
+    next slot's token — the average-case cost win the docstring above
+    describes. A row that never fires within `k_max` is capped there
+    (`K_MAX` stays a safety cap, not a modeling ceiling).

    Under `history="attention"` the history encoding is computed once per
-    slot via `Stage2Autoregressive.history_step` (a KV-cache append, §10)
+    slot via `Stage2Autoregressive.history_step` (a KV-cache append)
    rather than re-derived by every model call inside that slot — so an ODE
    loop's `steps` substeps, and the separate `predict_type` call when the
    type slice isn't folded into the trunk output, all reuse the SAME `hist`
@@ -291,8 +304,8 @@ def sample_secondaries_ar(
    device = cond_cont.device
    k_max = sec_decoder.k_max
    type_dim = sec_decoder.type_dim
-    generator = sec_decoder.generator_kind
-    target = sec_decoder.particle_type_cfg.get("target", "physical")
+    objective = build_objective(sec_decoder.generator_kind)
+    target = sec_decoder.particle_type_cfg.target
    type_folded = _type_folded(sec_decoder)
    token_dim = CONT_SLOT_DIM + type_dim if type_folded else CONT_SLOT_DIM

@@ -304,18 +317,44 @@ def sample_secondaries_ar(
    remaining = torch.ones(B, device=device)
    history_cache = sec_decoder.init_history_cache()

+    use_stop_token = n_sec_pred is None
+    if use_stop_token:
+        assert getattr(sec_decoder, "stop_head", None) is not None, (
+            "sample_secondaries_ar called with n_sec_pred=None on a decoder "
+            "with no stop_head — only valid under stage2_model.n_sec.mode = "
+            "'stop_token'"
+        )
+        finished = torch.zeros(B, dtype=torch.bool, device=device)
+        derived_n_sec = torch.full((B,), k_max, dtype=torch.long, device=device)
+
    for k in range(k_max):
        has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device)
        history_feat = prev_repr.unsqueeze(1)  # (B, 1, CONT_SLOT_DIM + type_dim)
        remaining_frac = remaining.unsqueeze(1)  # (B, 1)
-        slot_idx = torch.full(
-            (B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32
-        )
-        hist, history_cache = sec_decoder.history_step(
-            history_feat, has_prev, history_cache
-        )
+        slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
+        hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache)

-        if generator == "wgan":
+        if use_stop_token:
+            stop_logit = sec_decoder.predict_stop(
+                cond_cont,
+                cond_cat,
+                stage1_out,
+                history_feat,
+                has_prev,
+                remaining_frac,
+                slot_idx,
+                hist=hist,
+            ).squeeze(1)
+            if sec_decoder.stop_sampling == "sample":
+                stop_now = torch.rand(B, device=device) < torch.sigmoid(stop_logit)
+            else:
+                stop_now = stop_logit >= 0.0
+            derived_n_sec[stop_now & ~finished] = k
+            finished = finished | stop_now
+            if finished.all():
+                break
+
+        if objective.is_adversarial:
            z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
            token = sec_decoder(
                z,
@@ -368,30 +407,25 @@ def sample_secondaries_ar(
        sec_type[:, k] = type_k

        if target == "onehot":
-            type_for_history = F.one_hot(
-                type_k.argmax(dim=-1), num_classes=type_dim
-            ).float()
+            type_for_history = F.one_hot(type_k.argmax(dim=-1), num_classes=type_dim).float()
        else:
            type_for_history = type_k

        stick_fraction = torch.sigmoid(cont_k[:, 0])
-        prev_repr = torch.cat(
-            [stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1
-        )
+        prev_repr = torch.cat([stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1)
        remaining = torch.clamp(remaining * (1.0 - stick_fraction), min=0.0)

-    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(
-        1
-    )
+    resolved_n_sec = derived_n_sec if use_stop_token else n_sec_pred
+    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < resolved_n_sec.unsqueeze(1)
    return sec_cont, sec_type, sec_valid


 # ---------------------------------------------------------------------------
 # Per-stage dispatch — shared by giant/rollout.py and giant/cli.py's
 # `predict` command, since both need "given a stage model, produce a
-# sample" without hand-picking the sampler themselves (docs/v0.3.0-design.md
-# decision 2: each stage's generative objective is independent, read off the
-# model's own `generator_kind`, not a caller-supplied `mode` string).
+# sample" without hand-picking the sampler themselves (each stage's
+# generative objective is independent, read off the model's own
+# `generator_kind`, not a caller-supplied `mode` string).
 # ---------------------------------------------------------------------------


@@ -403,10 +437,10 @@ def sample_stage1(
    ddpm_steps: int = 1000,
 ) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Dispatches on `stage1_model.generator_kind`."""
-    kind = stage1_model.generator_kind
-    if kind == "wgan":
+    objective = build_objective(stage1_model.generator_kind)
+    if objective.is_adversarial:
        return sample_wgan(stage1_model, cond_cont, cond_cat)
-    if kind == "ddpm":
+    if isinstance(objective, DdpmObjective):
        schedule = CosineSchedule(T=ddpm_steps).to(cond_cont.device)
        return sample_ddpm(stage1_model, cond_cont, cond_cat, schedule)
    return sample_flow(stage1_model, cond_cont, cond_cat, steps=steps)
@@ -417,7 +451,7 @@ def sample_stage2(
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    stage1_out: torch.Tensor,
-    n_sec_pred: torch.Tensor,
+    n_sec_pred: torch.Tensor | None,
    steps: int,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Dispatches on `decoder` (one-shot vs autoregressive — the class
@@ -426,18 +460,18 @@ def sample_stage2(
    built with `generator="ddpm"` in practice and `flow_matching_loss_secondary*`
    is the only stage-2 training path that exists for the non-adversarial
    case, so there's nothing to dispatch to here.
+
+    `n_sec_pred=None` (from `resolve_n_sec` on a stop-token decoder) is only
+    meaningful for the autoregressive path — see `sample_secondaries_ar`'s
+    docstring; the one-shot samplers have no per-token stop mechanism to
+    derive a count from, so `n_sec_pred` must already be resolved for them.
    """
    if isinstance(sec_decoder, Stage2Autoregressive):
-        return sample_secondaries_ar(
-            sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps
-        )
-    if sec_decoder.generator_kind == "wgan":
-        return sample_secondaries_wgan(
-            sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred
-        )
-    return sample_secondaries(
-        sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps
-    )
+        return sample_secondaries_ar(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
+    assert n_sec_pred is not None, "one-shot stage-2 decoders need a resolved n_sec_pred"
+    if build_objective(sec_decoder.generator_kind).is_adversarial:
+        return sample_secondaries_wgan(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred)
+    return sample_secondaries(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)


 def resolve_n_sec(
@@ -447,22 +481,31 @@ def resolve_n_sec(
    cond_cat: torch.Tensor,
    stage1_out: torch.Tensor,
    n_sec_pred: torch.Tensor | None,
-) -> torch.Tensor:
+) -> torch.Tensor | None:
    """`n_sec_pred` is already populated when `stage1_model` owns a legacy
    `n_sec_head` (a migrated v0.2 checkpoint — see `Stage1Model`'s
-    docstring); otherwise ask stage 2, which owns it by default under
-    decision 1 (docs/v0.3.0-design.md §2). Raises if neither stage owns a
-    head at all — the only way that happens is `stage2_model.n_sec.mode`
-    other than `"head"` (`"truth"`/`"stop_token"`), neither of which is a
-    valid rollout-/predict-capable checkpoint (§3.3, §9)."""
+    docstring); otherwise ask stage 2, which owns it by default.
+
+    Returns `None` when `sec_decoder` owns a `stop_head` (`n_sec.mode =
+    "stop_token"`) instead of an `n_sec_head` — there is nothing to resolve
+    up front in that case, since the count only exists once
+    `sample_secondaries_ar` has actually generated (or stopped generating)
+    tokens; the caller passes this `None` straight through to `sample_stage2`
+    and reads the real count back off its returned `sec_valid`
+    (`sec_valid.sum(-1)`) afterwards.
+
+    Raises if neither stage owns any n_sec mechanism at all — the only way
+    that happens is `stage2_model.n_sec.mode = "truth"`, which is not a valid
+    rollout-/predict-capable checkpoint."""
    if n_sec_pred is not None:
        return n_sec_pred
+    if getattr(sec_decoder, "stop_head", None) is not None:
+        return None
    if getattr(sec_decoder, "n_sec_head", None) is None:
        raise RuntimeError(
-            "checkpoint has no n_sec_head on either stage — needs "
-            "stage2_model.n_sec.mode = 'head' (the default); 'truth' is "
-            "standalone-evaluation-only and 'stop_token' isn't implemented "
-            "(docs/v0.3.0-design.md §3.3/§9)"
+            "checkpoint has no n_sec_head/stop_head on either stage — needs "
+            "stage2_model.n_sec.mode = 'head' (the default) or 'stop_token'; "
+            "'truth' is standalone-evaluation-only"
        )
    logits = sec_decoder.predict_n_sec(cond_cont, cond_cat, stage1_out)
    return logits.argmax(dim=-1)
@@ -1,5 +1,5 @@
 """Cut a new raw generation or processed schema version for the geant_steps
-dataset tree (see scripts/migrate_geant_steps.py for the layout):
+dataset tree (see giant/tools/migrate_geant_steps.py for the layout):

    raw/<kind>/<gen>/<detector>/shard-NNN.root
    processed/<kind>/<gen>/<schema>/<detector>/shard-NNN.parquet
@@ -82,9 +82,7 @@ def _max_index(parent: Path, pattern: re.Pattern) -> int:

 def _git_user_name() -> str | None:
    try:
-        out = subprocess.run(
-            ["git", "config", "user.name"], capture_output=True, text=True, timeout=2
-        )
+        out = subprocess.run(["git", "config", "user.name"], capture_output=True, text=True, timeout=2)
    except (OSError, subprocess.SubprocessError):
        # OSError (e.g. git not on PATH) and subprocess.SubprocessError
        # (e.g. TimeoutExpired) are unrelated hierarchies — TimeoutExpired
@@ -142,9 +140,7 @@ def plan_bump_schema(
    raw_gen_dir = root / "raw" / kind / gen_tag
    processed_gen_dir = root / "processed" / kind / gen_tag
    if not raw_gen_dir.is_dir() and not processed_gen_dir.is_dir():
-        raise SystemExit(
-            f"error: {gen_tag} doesn't exist yet for kind={kind} — run bump-gen first"
-        )
+        raise SystemExit(f"error: {gen_tag} doesn't exist yet for kind={kind} — run bump-gen first")
    if target is not None:
        if not SCHEMA_RE.match(target):
            raise SystemExit(f"error: --to must look like 'schemaN', got {target!r}")
@@ -154,9 +150,7 @@ def plan_bump_schema(
        schema_tag = f"schema{next_schema}"
    new_dirs = [processed_gen_dir / schema_tag]
    by_suffix = f" ({by})" if by else ""
-    log_line = (
-        f"- `{gen_tag}`/`{schema_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
-    )
+    log_line = f"- `{gen_tag}`/`{schema_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
    return new_dirs, log_line


@@ -212,9 +206,7 @@ def _manifest_referenced_files(pools_root: Path) -> set[Path]:
    return referenced


-def _referenced_root_count(
-    raw_gen_dir: Path, processed_gen_dir: Path
-) -> tuple[int, int]:
+def _referenced_root_count(raw_gen_dir: Path, processed_gen_dir: Path) -> tuple[int, int]:
    """(total .root files, count with a same-named .parquet under any schema) for one gen."""
    if not raw_gen_dir.is_dir():
        return 0, 0
@@ -243,9 +235,7 @@ def _referenced_root_count(
    return total, referenced


-def _referenced_parquet_count(
-    schema_dir: Path, manifest_referenced: set[Path]
-) -> tuple[int, int]:
+def _referenced_parquet_count(schema_dir: Path, manifest_referenced: set[Path]) -> tuple[int, int]:
    """(total .parquet files, count listed in at least one manifest) for one schema dir."""
    if not schema_dir.is_dir():
        return 0, 0
@@ -342,11 +332,7 @@ def print_status(root: Path) -> None:
    grand_files = 0
    for kind_dir in sorted(p for p in raw_root.iterdir() if p.is_dir()):
        kind = kind_dir.name
-        gens = sorted(
-            int(m.group(1))
-            for m in (GEN_RE.match(p.name) for p in kind_dir.iterdir() if p.is_dir())
-            if m
-        )
+        gens = sorted(int(m.group(1)) for m in (GEN_RE.match(p.name) for p in kind_dir.iterdir() if p.is_dir()) if m)
        print(_colorize(f"{kind}/", "kind"))
        kind_total = 0
        kind_files = 0
@@ -359,25 +345,18 @@ def print_status(root: Path) -> None:
            schemas = sorted(
                int(m.group(1))
                for m in (
-                    SCHEMA_RE.match(p.name)
-                    for p in (schema_dir.iterdir() if schema_dir.is_dir() else [])
-                    if p.is_dir()
+                    SCHEMA_RE.match(p.name) for p in (schema_dir.iterdir() if schema_dir.is_dir() else []) if p.is_dir()
                )
                if m
            )
            schema_sizes = {s: _du(schema_dir / f"schema{s}") for s in schemas}
            schema_counts = {
-                s: _referenced_parquet_count(
-                    schema_dir / f"schema{s}", manifest_referenced
-                )
-                for s in schemas
+                s: _referenced_parquet_count(schema_dir / f"schema{s}", manifest_referenced) for s in schemas
            }
            processed_size = sum(schema_sizes.values())
            processed_files = sum(c[0] for c in schema_counts.values())
            processed_referenced = sum(c[1] for c in schema_counts.values())
-            raw_files, raw_referenced = _referenced_root_count(
-                raw_gen_dir, processed_gen_dir
-            )
+            raw_files, raw_referenced = _referenced_root_count(raw_gen_dir, processed_gen_dir)
            gen_total = raw_size + processed_size
            gen_files = raw_files + processed_files
            kind_total += gen_total
@@ -425,9 +404,7 @@ def print_status(root: Path) -> None:
                        print(_reason_line(schema_reason, indent=4))
            else:
                print(_colorize("      (none)", "schema"))
-        print(
-            _row(f"{kind} total", kind_total, indent=1, level="gen", count=kind_files)
-        )
+        print(_row(f"{kind} total", kind_total, indent=1, level="gen", count=kind_files))
        print()
        grand_total += kind_total
        grand_files += kind_files
@@ -526,13 +503,8 @@ def plan_update_manifest(
    return result, missing


-def apply_update_manifest(
-    manifest_path: Path, lines: list[tuple[str, str | None]]
-) -> None:
-    out = [
-        replacement if replacement is not None else original
-        for original, replacement in lines
-    ]
+def apply_update_manifest(manifest_path: Path, lines: list[tuple[str, str | None]]) -> None:
+    out = [replacement if replacement is not None else original for original, replacement in lines]
    manifest_path.write_text("\n".join(out) + "\n")


@@ -552,9 +524,7 @@ def _resolve_manifest_files(manifest_path: Path) -> list[Path]:
    return files


-def plan_create_manifest(
-    output_path: Path, parquet_files: list[Path]
-) -> tuple[list[str], list[Path], list[Path]]:
+def plan_create_manifest(output_path: Path, parquet_files: list[Path]) -> tuple[list[str], list[Path], list[Path]]:
    """Return (relative_lines, missing_files, resolved_abs_paths)."""
    manifest_dir = output_path.resolve().parent
    lines: list[str] = []
@@ -569,9 +539,7 @@ def plan_create_manifest(
    return lines, missing, resolved


-def check_holdout_overlap(
-    output_path: Path, resolved_new_files: list[Path]
-) -> list[tuple[str, Path]]:
+def check_holdout_overlap(output_path: Path, resolved_new_files: list[Path]) -> list[tuple[str, Path]]:
    """Return (other_manifest_name, file) pairs where new files clash with existing manifests.

    The check is triggered when output_path is (or will be) holdout.manifest, or when a
@@ -611,7 +579,7 @@ def apply_create_manifest(output_path: Path, lines: list[str]) -> None:


 # ---------------------------------------------------------------------------
-# CLI entry points (called from scripts/dwarf.py)
+# CLI entry points (called from giant/tools/dwarf.py)
 # ---------------------------------------------------------------------------


@@ -641,9 +609,7 @@ def _run_bump(
    if gen is None:
        new_dirs, log_line = plan_bump_gen(root_path, kind, reason, by, date, to)
    else:
-        new_dirs, log_line = plan_bump_schema(
-            root_path, kind, gen, reason, by, date, to
-        )
+        new_dirs, log_line = plan_bump_schema(root_path, kind, gen, reason, by, date, to)

    print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
    print("new directories:")
@@ -1,6 +1,5 @@
-"""Portal-machine follow-up for v0.3.0 step 2 (docs/v0.3.0-design.md §4.3):
-diff a real v0.2 checkpoint's outputs against the new `build_models` on the
-same input batch.
+"""Portal-machine follow-up for v0.3.0 step 2: diff a real v0.2 checkpoint's
+outputs against the new `build_models` on the same input batch.

 `tests/test_migration_v02_v03.py` already proves this bit-identical with
 synthetic random weights, but that test can't run where it matters (no
@@ -11,12 +10,12 @@ machine against an actual trained checkpoint before merging

 Usage (from the repo root, on a portal machine):

-    uv run python scripts/check_migration_v02_v03.py /ceph/lbogner/.../best.pt
-    uv run python scripts/check_migration_v02_v03.py /ceph/lbogner/.../best.pt --ema
-    uv run python scripts/check_migration_v02_v03.py /ceph/lbogner/.../best.pt --batch 32 --seed 1
+    uv run python giant/tools/check_migration_v02_v03.py /ceph/lbogner/.../best.pt
+    uv run python giant/tools/check_migration_v02_v03.py /ceph/lbogner/.../best.pt --ema
+    uv run python giant/tools/check_migration_v02_v03.py /ceph/lbogner/.../best.pt --batch 32 --seed 1

 Run it once against a flow (or ddpm) checkpoint and once against a wgan
-checkpoint (design doc §4.3's "one flow checkpoint and one WGAN checkpoint").
+checkpoint ("one flow checkpoint and one WGAN checkpoint").
 A routed checkpoint (`model_config["router"]["enabled"]`) is only checked for
 successful construction — `giant.model.network.migrate_legacy_state_dict`
 doesn't yet remap routed (Expert-per-router) state dicts, so the
@@ -85,19 +84,12 @@ def main() -> int:
    mode = model_config.get("mode", "flow")
    routed = bool((model_config.get("router") or {}).get("enabled"))
    print(f"checkpoint: {args.checkpoint}")
-    print(
-        f"  mode={mode!r}  conditioning={model_config.get('conditioning')!r}  "
-        f"routed={routed}  ema={args.ema}"
-    )
+    print(f"  mode={mode!r}  conditioning={model_config.get('conditioning')!r}  routed={routed}  ema={args.ema}")

    stage1_key = "model_ema" if args.ema and "model_ema" in ckpt else "model"
-    stage2_key = (
-        "sec_decoder_ema" if args.ema and "sec_decoder_ema" in ckpt else "sec_decoder"
-    )
+    stage2_key = "sec_decoder_ema" if args.ema and "sec_decoder_ema" in ckpt else "sec_decoder"
    if args.ema and stage1_key == "model":
-        print(
-            "  warning: --ema requested but no model_ema in checkpoint, using raw weights"
-        )
+        print("  warning: --ema requested but no model_ema in checkpoint, using raw weights")

    # --- old side: the frozen v0.2 snapshot, loaded with the checkpoint's own weights ---
    old_stage1, old_stage2 = legacy.build_models(model_config)
@@ -115,15 +107,12 @@ def main() -> int:
        print(
            "  routed checkpoint: migrate_legacy_state_dict only handles the "
            "monolithic trunk shape — verifying construction only, skipping "
-            "the bit-identical weight/output comparison. See "
-            "docs/v0.3.0-design.md §2.4's scope note."
+            "the bit-identical weight/output comparison."
        )
        print("PASS (construction only, routed checkpoint)")
        return 0

-    remapped1, remapped2 = net.migrate_legacy_state_dict(
-        ckpt[stage1_key], ckpt[stage2_key]
-    )
+    remapped1, remapped2 = net.migrate_legacy_state_dict(ckpt[stage1_key], ckpt[stage2_key])
    missing1, unexpected1 = new_stage1.load_state_dict(remapped1, strict=True)
    missing2, unexpected2 = new_stage2.load_state_dict(remapped2, strict=True)
    if missing1 or unexpected1 or missing2 or unexpected2:
@@ -134,9 +123,7 @@ def main() -> int:
    new_stage1.eval()
    new_stage2.eval()

-    cond_cont, cond_cat, x1, x2, t, z1, z2 = _random_batch(
-        model_config, args.batch, args.seed
-    )
+    cond_cont, cond_cat, x1, x2, t, z1, z2 = _random_batch(model_config, args.batch, args.seed)

    ok = True
    with torch.no_grad():
@@ -32,7 +32,7 @@ from dataclasses import dataclass
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path

-# Must match scripts/bump_dataset_version.py's GEN_RE.
+# Must match giant/tools/bump_dataset_version.py's GEN_RE.
 GEN_RE = re.compile(r"^gen\d+$")
 SHARD_RE = re.compile(r"^shard-(\d+)\.root$")

@@ -64,9 +64,7 @@ def parse_detector_spec(spec: str) -> tuple[str, str | None]:
    if ":" in spec:
        label, config = spec.split(":", 1)
        if not label or not config:
-            raise PlanError(
-                f"invalid --detector spec {spec!r}: expected NAME or NAME:CONFIG"
-            )
+            raise PlanError(f"invalid --detector spec {spec!r}: expected NAME or NAME:CONFIG")
        return label, config
    return spec, None

@@ -94,9 +92,7 @@ def plan_jobs(
        raise PlanError(f"--gen must look like 'genN', got {gen!r}")
    gen_dir = dataset_root / "raw" / kind / gen
    if not gen_dir.is_dir():
-        raise PlanError(
-            f"{gen_dir} doesn't exist — run bump_dataset_version.py bump-gen first"
-        )
+        raise PlanError(f"{gen_dir} doesn't exist — run bump_dataset_version.py bump-gen first")

    jobs = []
    for spec in detector_specs:
@@ -124,9 +120,7 @@ def job_seed(kind: str, gen: str, job: SimJob, energy_gev: float | None) -> int:
    return zlib.crc32(key.encode()) & 0x7FFFFFFF


-def build_cmd(
-    executable: Path, job: SimJob, events_per_file: int, energy_gev: float | None
-) -> list[str]:
+def build_cmd(executable: Path, job: SimJob, events_per_file: int, energy_gev: float | None) -> list[str]:
    """minicalosim executables take positional `[configName] nEvents [energy_GeV]`."""
    cmd = [str(executable)]
    if job.config:
@@ -147,10 +141,7 @@ def run_job(
    gen: str,
    tmp_root: Path,
 ) -> JobResult:
-    workdir = (
-        tmp_root
-        / f"{kind}-{gen}-{job.detector}-{job.shard_index:03d}-{uuid.uuid4().hex[:8]}"
-    )
+    workdir = tmp_root / f"{kind}-{gen}-{job.detector}-{job.shard_index:03d}-{uuid.uuid4().hex[:8]}"
    workdir.mkdir(parents=True)

    cmd = build_cmd(executable, job, events_per_file, energy_gev)
@@ -174,20 +165,12 @@ def run_job(
            job,
            False,
            None,
-            f"expected exactly one .root output in {workdir}, found {len(produced)}: "
-            f"{[p.name for p in produced]}",
+            f"expected exactly one .root output in {workdir}, found {len(produced)}: {[p.name for p in produced]}",
            result.stdout,
            result.stderr,
        )

-    dest = (
-        dataset_root
-        / "raw"
-        / kind
-        / gen
-        / job.detector
-        / f"shard-{job.shard_index:03d}.root"
-    )
+    dest = dataset_root / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root"
    if dest.exists():
        return JobResult(
            job,
@@ -279,14 +262,7 @@ def run_make_root(
    print(f"executable: {executable}")
    for job in planned_jobs:
        cmd = build_cmd(executable, job, events_per_file, energy_gev)
-        dest = (
-            dataset_root_path
-            / "raw"
-            / kind
-            / gen
-            / job.detector
-            / f"shard-{job.shard_index:03d}.root"
-        )
+        dest = dataset_root_path / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root"
        seed = job_seed(kind, gen, job, energy_gev)
        print(f"  MINICALOSIM_SEED={seed} {' '.join(cmd)}  ->  {dest}")

@@ -1,6 +1,6 @@
 """dwarf — little helper to `giant`: dataset/tooling CLI for the geant_steps pipeline.

-Unifies the standalone scripts/*.py conversion, migration, versioning, and
+Unifies the standalone giant/tools/*.py conversion, migration, versioning, and
 simulation-fanout tools into one Typer app so there's a single command name
 (and `--help`) to remember instead of five differently-hyphenated ones.
 """
@@ -14,20 +14,20 @@ import typer
 from typing_extensions import Annotated

 from giant.config import Conditioning
-from scripts.bump_dataset_version import (
+from giant.tools.bump_dataset_version import (
    run_bump_gen,
    run_bump_schema,
    run_create_manifest,
    run_status,
    run_update_manifest,
 )
-from scripts.create_root_files import run_make_root
-from scripts.geometry_oracle import run_build_geometry_oracle
-from scripts.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
-from scripts.migrate_geant_steps import run_migration
-from scripts.steps_to_parquet import convert_steps_to_parquet
-from scripts.steps_to_parquet_parallel import run_parallel_job
-from scripts.warm_setup_cache import run_warm_setup_cache
+from giant.tools.create_root_files import run_make_root
+from giant.tools.geometry_oracle import run_build_geometry_oracle
+from giant.tools.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
+from giant.tools.migrate_geant_steps import run_migration
+from giant.tools.steps_to_parquet import convert_steps_to_parquet
+from giant.tools.steps_to_parquet_parallel import run_parallel_job
+from giant.tools.warm_setup_cache import run_warm_setup_cache

 app = typer.Typer(no_args_is_help=True)

@@ -80,8 +80,7 @@ def convert(
        typer.Option(
            "--output",
            "-o",
-            help="Output Parquet file (default: <input>.parquet). Only valid "
-            "with a single input file and --jobs 1.",
+            help="Output Parquet file (default: <input>.parquet). Only valid with a single input file and --jobs 1.",
        ),
    ] = None,
    batch_size: Annotated[
@@ -91,9 +90,7 @@ def convert(
            help="Uproot read batch size, e.g. '100 MB' or '500000' (rows)",
        ),
    ] = "100 MB",
-    tree: Annotated[
-        str, typer.Option("--tree", help="Tree name inside the ROOT file")
-    ] = "Steps",
+    tree: Annotated[str, typer.Option("--tree", help="Tree name inside the ROOT file")] = "Steps",
    compression: Annotated[
        Compression, typer.Option("--compression", help="Parquet compression codec")
    ] = Compression.snappy,
@@ -129,15 +126,11 @@ def convert(
        raise typer.Exit(1)
    _warn_if_exceeds_shared_quota(jobs, "--jobs")

-    compression_value = (
-        "uncompressed" if compression is Compression.none else compression.value
-    )
+    compression_value = "uncompressed" if compression is Compression.none else compression.value

    if jobs == 1:
        if output is not None and len(root_files) > 1:
-            typer.echo(
-                "error: --output can only be used with a single input file", err=True
-            )
+            typer.echo("error: --output can only be used with a single input file", err=True)
            raise typer.Exit(1)
        total_orphaned = 0
        for root_file in root_files:
@@ -150,10 +143,7 @@ def convert(
            )
            total_orphaned += n_orphaned
        if total_orphaned:
-            typer.echo(
-                f"\n{total_orphaned} orphaned child track(s) dropped across "
-                f"{len(root_files)} file(s)."
-            )
+            typer.echo(f"\n{total_orphaned} orphaned child track(s) dropped across {len(root_files)} file(s).")
        return

    if output is not None:
@@ -176,9 +166,7 @@ def convert(

@app.command()
 def migrate(
-    root: Annotated[
-        Path, typer.Argument(help="Dataset root to migrate in place")
-    ] = _DATASET_ROOT_DEFAULT,
+    root: Annotated[Path, typer.Argument(help="Dataset root to migrate in place")] = _DATASET_ROOT_DEFAULT,
    execute: Annotated[
        bool,
        typer.Option(
@@ -190,8 +178,7 @@ def migrate(
        bool,
        typer.Option(
            "--copy",
-            help="Copy instead of move, leaving the originals in place "
-            "(e.g. if another process is still reading them)",
+            help="Copy instead of move, leaving the originals in place (e.g. if another process is still reading them)",
        ),
    ] = False,
 ) -> None:
@@ -202,12 +189,8 @@ def migrate(
@app.command("bump-gen")
 def bump_gen(
    reason: Annotated[str, typer.Option("--reason", help="Why this gen exists")],
-    kind: Annotated[
-        str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
-    ] = "steps",
-    by: Annotated[
-        Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")
-    ] = None,
+    kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
+    by: Annotated[Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")] = None,
    date: Annotated[
        Optional[str],
        typer.Option("--date", help="Override date (default: today, ISO)"),
@@ -220,12 +203,8 @@ def bump_gen(
            help="Target gen tag (default: one past the current highest)",
        ),
    ] = None,
-    execute: Annotated[
-        bool, typer.Option("--execute", help="Apply (default: dry run)")
-    ] = False,
-    root: Annotated[
-        Path, typer.Option("--root", help="Dataset root")
-    ] = _DATASET_ROOT_DEFAULT,
+    execute: Annotated[bool, typer.Option("--execute", help="Apply (default: dry run)")] = False,
+    root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
    """Cut a new raw generation."""
    run_bump_gen(
@@ -243,12 +222,8 @@ def bump_gen(
 def bump_schema(
    gen: Annotated[str, typer.Option("--gen", help="Existing gen tag, e.g. gen1")],
    reason: Annotated[str, typer.Option("--reason", help="Why this schema exists")],
-    kind: Annotated[
-        str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
-    ] = "steps",
-    by: Annotated[
-        Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")
-    ] = None,
+    kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
+    by: Annotated[Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")] = None,
    date: Annotated[
        Optional[str],
        typer.Option("--date", help="Override date (default: today, ISO)"),
@@ -261,12 +236,8 @@ def bump_schema(
            help="Target schema tag (default: one past the current highest)",
        ),
    ] = None,
-    execute: Annotated[
-        bool, typer.Option("--execute", help="Apply (default: dry run)")
-    ] = False,
-    root: Annotated[
-        Path, typer.Option("--root", help="Dataset root")
-    ] = _DATASET_ROOT_DEFAULT,
+    execute: Annotated[bool, typer.Option("--execute", help="Apply (default: dry run)")] = False,
+    root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
    """Cut a new schema within a gen."""
    run_bump_schema(
@@ -283,9 +254,7 @@ def bump_schema(

@app.command()
 def status(
-    root: Annotated[
-        Path, typer.Option("--root", help="Dataset root")
-    ] = _DATASET_ROOT_DEFAULT,
+    root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
    """List existing gens/schemas per kind."""
    run_status(str(root))
@@ -293,9 +262,7 @@ def status(

@app.command("update-manifest")
 def update_manifest(
-    manifests: Annotated[
-        list[Path], typer.Argument(help="One or more .manifest files to update")
-    ],
+    manifests: Annotated[list[Path], typer.Argument(help="One or more .manifest files to update")],
    schema: Annotated[
        Optional[str],
        typer.Option(
@@ -306,9 +273,7 @@ def update_manifest(
    ] = None,
    gen: Annotated[
        Optional[str],
-        typer.Option(
-            "--gen", metavar="genN", help="Target gen tag (default: keep existing gen)"
-        ),
+        typer.Option("--gen", metavar="genN", help="Target gen tag (default: keep existing gen)"),
    ] = None,
    execute: Annotated[
        bool,
@@ -316,9 +281,7 @@ def update_manifest(
    ] = False,
 ) -> None:
    """Repoint manifest(s) to a new gen and/or schema, verifying all target files exist."""
-    run_update_manifest(
-        [str(m) for m in manifests], schema=schema, execute=execute, gen=gen
-    )
+    run_update_manifest([str(m) for m in manifests], schema=schema, execute=execute, gen=gen)


@app.command("create-manifest")
@@ -333,22 +296,15 @@ def create_manifest(
        typer.Option(
            "--pool",
            metavar="DETECTOR",
-            help="Detector name; combined with --type and --root to form "
-            "<root>/pools/<detector>/<type>.manifest",
+            help="Detector name; combined with --type and --root to form <root>/pools/<detector>/<type>.manifest",
        ),
    ] = None,
    type_: Annotated[
        Optional[PoolType],
-        typer.Option(
-            "--type", help="Pool type — full, holdout, or dev (required with --pool)"
-        ),
+        typer.Option("--type", help="Pool type — full, holdout, or dev (required with --pool)"),
    ] = None,
-    root: Annotated[
-        Path, typer.Option("--root", help="Dataset root (used with --pool)")
-    ] = _DATASET_ROOT_DEFAULT,
-    execute: Annotated[
-        bool, typer.Option("--execute", help="Write the manifest (default: dry run)")
-    ] = False,
+    root: Annotated[Path, typer.Option("--root", help="Dataset root (used with --pool)")] = _DATASET_ROOT_DEFAULT,
+    execute: Annotated[bool, typer.Option("--execute", help="Write the manifest (default: dry run)")] = False,
    force: Annotated[
        bool,
        typer.Option("--force", help="Overwrite the manifest if it already exists"),
@@ -368,9 +324,7 @@ def create_manifest(

@app.command("make-root")
 def make_root(
-    executable: Annotated[
-        Path, typer.Option("--executable", help="Built minicalosim run_* executable")
-    ],
+    executable: Annotated[Path, typer.Option("--executable", help="Built minicalosim run_* executable")],
    detector: Annotated[
        list[str],
        typer.Option(
@@ -382,15 +336,9 @@ def make_root(
            "Repeatable.",
        ),
    ],
-    num_files: Annotated[
-        int, typer.Option("--num-files", help="New shards to create per detector")
-    ],
-    events_per_file: Annotated[
-        int, typer.Option("--events-per-file", help="nEvents passed to the executable")
-    ],
-    gen: Annotated[
-        str, typer.Option("--gen", help="Existing gen tag under raw/<kind>/, e.g. gen1")
-    ],
+    num_files: Annotated[int, typer.Option("--num-files", help="New shards to create per detector")],
+    events_per_file: Annotated[int, typer.Option("--events-per-file", help="nEvents passed to the executable")],
+    gen: Annotated[str, typer.Option("--gen", help="Existing gen tag under raw/<kind>/, e.g. gen1")],
    energy_gev: Annotated[
        float | None,
        typer.Option(
@@ -401,20 +349,12 @@ def make_root(
            "to name the dataset accordingly.",
        ),
    ] = None,
-    kind: Annotated[
-        str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
-    ] = "steps",
-    dataset_root: Annotated[
-        Path, typer.Option("--dataset-root", help="Dataset root")
-    ] = _DATASET_ROOT_DEFAULT,
-    jobs: Annotated[
-        int, typer.Option("--jobs", "-j", help="Parallel simulation runs (default: 4)")
-    ] = 4,
+    kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
+    dataset_root: Annotated[Path, typer.Option("--dataset-root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
+    jobs: Annotated[int, typer.Option("--jobs", "-j", help="Parallel simulation runs (default: 4)")] = 4,
    execute: Annotated[
        bool,
-        typer.Option(
-            "--execute", help="Actually run jobs (default: dry run / print plan)"
-        ),
+        typer.Option("--execute", help="Actually run jobs (default: dry run / print plan)"),
    ] = False,
 ) -> None:
    """Generate new ROOT shards via a minicalosim executable."""
@@ -441,9 +381,7 @@ class OracleMethod(str, Enum):

@app.command("build-geometry-oracle")
 def build_geometry_oracle(
-    data: Annotated[
-        Path, typer.Argument(help="Steps parquet file or directory of steps files")
-    ],
+    data: Annotated[Path, typer.Argument(help="Steps parquet file or directory of steps files")],
    out: Annotated[Path, typer.Option("--out", "-o", help="Output oracle .pkl path")],
    method: Annotated[
        OracleMethod,
@@ -456,9 +394,7 @@ def build_geometry_oracle(
            ),
        ),
    ] = OracleMethod.slab,
-    k: Annotated[
-        int, typer.Option("--k", help="Neighbours for the knn classifier")
-    ] = 1,
+    k: Annotated[int, typer.Option("--k", help="Neighbours for the knn classifier")] = 1,
    subsample: Annotated[
        int,
        typer.Option("--subsample", help="Max reference points sampled from the data"),
@@ -504,76 +440,107 @@ def build_geometry_oracle(
 def warm_cache(
    data: Annotated[
        Path,
-        typer.Argument(
-            help="Parquet file, directory, or .manifest — same as `giant train`'s"
-        ),
+        typer.Argument(help="Parquet file, directory, or .manifest — same as `giant train`'s"),
    ],
+    config: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--config",
+            "-c",
+            help="TOML config file to warm for — same file the `giant train` run(s) will use. "
+            "Mutually exclusive with the flags below (put val-fraction/seed/conditioning/router "
+            "settings in the file itself, so warming and training can't disagree on them)",
+        ),
+    ] = None,
    val_fraction: Annotated[
-        float,
+        Optional[float],
        typer.Option(
            "--val-fraction",
            "-f",
-            help="Must match the `giant train` run(s) to warm for",
+            help="Must match the `giant train` run(s) to warm for. Not allowed together with --config",
        ),
-    ] = 0.1,
+    ] = None,
    seed: Annotated[
-        int,
+        Optional[int],
        typer.Option(
-            "--seed", "-s", help="Must match the `giant train` run(s) to warm for"
+            "--seed",
+            "-s",
+            help="Must match the `giant train` run(s) to warm for. Not allowed together with --config",
        ),
-    ] = 0,
+    ] = None,
    particle_conditioning: Annotated[
-        Conditioning,
+        Optional[Conditioning],
        typer.Option(
            "--particle-conditioning",
-            help="Must match the `giant train` run(s)' conditioning.particle.type "
-            "to warm for",
+            help="Must match the `giant train` run(s)' conditioning.particle.type to warm for. "
+            "Not allowed together with --config",
        ),
-    ] = Conditioning.physical,
+    ] = None,
    material_conditioning: Annotated[
-        Conditioning,
+        Optional[Conditioning],
        typer.Option(
            "--material-conditioning",
            help="Must match the `giant train` run(s)' conditioning.material.type "
            "to warm for — independent of --particle-conditioning "
-            "(docs/v0.3.0-design.md §3.1: the two axes may differ)",
+            "(the two axes may differ). Not allowed together with --config",
        ),
-    ] = Conditioning.physical,
+    ] = None,
    router: Annotated[
-        bool,
+        Optional[bool],
        typer.Option(
            "--router/--no-router",
-            help="Warm the process vocabulary too (only takes effect with "
-            "--router-type process)",
+            help="Warm the process vocabulary too (only takes effect with --router-type process). "
+            "Not allowed together with --config",
        ),
-    ] = False,
+    ] = None,
    router_type: Annotated[
-        str, typer.Option("--router-type", help="Router implementation name")
-    ] = "energy",
+        Optional[str],
+        typer.Option("--router-type", help="Router implementation name. Not allowed together with --config"),
+    ] = None,
    n_experts: Annotated[
-        int, typer.Option("--n-experts", help="Number of routed experts")
-    ] = 4,
+        Optional[int],
+        typer.Option("--n-experts", help="Number of routed experts. Not allowed together with --config"),
+    ] = None,
    rebuild: Annotated[
        bool,
-        typer.Option(
-            "--rebuild", help="Ignore any existing sidecar and recompute every section"
-        ),
+        typer.Option("--rebuild", help="Ignore any existing sidecar and recompute every section"),
    ] = False,
 ) -> None:
    """Precompute `giant train`'s setup-stage sidecar for `data` ahead of time.

    Warms the vocab maps, event-id split index, and the normalizer entry for
-    the given --val-fraction/--seed/--particle-conditioning/
-    --material-conditioning, so a later `giant train` run (or a `dwarf
-    hparam-scan` sweep, which shares one such entry across every run) skips
-    straight to training. See giant/data/setup_cache.py.
+    either --config, or the given --val-fraction/--seed/
+    --particle-conditioning/--material-conditioning/--router* flags, so a
+    later `giant train` run (or a `dwarf hparam-scan` sweep, which shares one
+    such entry across every run) skips straight to training. See
+    giant/data/setup_cache.py.
    """
+    flag_overrides = {
+        "--val-fraction": val_fraction,
+        "--seed": seed,
+        "--particle-conditioning": particle_conditioning,
+        "--material-conditioning": material_conditioning,
+        "--router/--no-router": router,
+        "--router-type": router_type,
+        "--n-experts": n_experts,
+    }
+    if config is not None:
+        given = [name for name, value in flag_overrides.items() if value is not None]
+        if given:
+            typer.echo(
+                f"error: --config cannot be combined with {', '.join(given)} "
+                "— put these settings in the config file instead",
+                err=True,
+            )
+            raise typer.Exit(1)
+
    run_warm_setup_cache(
        data=str(data),
+        config_path=config,
        val_fraction=val_fraction,
        seed=seed,
-        particle_conditioning=particle_conditioning.value,
-        material_conditioning=material_conditioning.value,
+        particle_conditioning=particle_conditioning.value if particle_conditioning is not None else None,
+        material_conditioning=material_conditioning.value if material_conditioning is not None else None,
        router_enabled=router,
        router_type=router_type,
        n_experts=n_experts,
@@ -38,9 +38,7 @@ def run_build_geometry_oracle(
        n_bins=n_bins,
    )

-    print(
-        f"method: {method}  reference points: {oracle.metadata['n_reference_points']:,}"
-    )
+    print(f"method: {method}  reference points: {oracle.metadata['n_reference_points']:,}")
    print("classes (material, layer_id):")
    for material, layer_id in oracle.classes:
        print(f"  {material:<12} layer_id={layer_id}")
@@ -68,8 +68,8 @@ def final_metrics(metrics_path: Path) -> tuple[int, float, float]:
    with open(metrics_path, newline="") as f:
        rows = list(csv.DictReader(f))
    epochs_completed = int(rows[-1]["epoch"])
-    final_val_loss = float(rows[-1]["val_loss"])
-    best_val_loss = min(float(r["val_loss"]) for r in rows)
+    final_val_loss = float(rows[-1]["val/loss"])
+    best_val_loss = min(float(r["val/loss"]) for r in rows)
    return epochs_completed, final_val_loss, best_val_loss


@@ -154,9 +154,7 @@ def run_hparam_scan(
        wall_time_s = time.monotonic() - start

        if metrics_path.exists():
-            epochs_completed, final_val_loss, best_val_loss = final_metrics(
-                metrics_path
-            )
+            epochs_completed, final_val_loss, best_val_loss = final_metrics(metrics_path)
            append_summary(
                summary_path,
                {
@@ -171,11 +169,6 @@ def run_hparam_scan(
                    "wall_time_s": round(wall_time_s, 1),
                },
            )
-            print(
-                f"[{i}/{len(runs)}] {name} — val_loss {final_val_loss:.4f} "
-                f"({wall_time_s:.1f}s)"
-            )
+            print(f"[{i}/{len(runs)}] {name} — val_loss {final_val_loss:.4f} ({wall_time_s:.1f}s)")
        else:
-            print(
-                f"[{i}/{len(runs)}] {name} — no metrics.csv produced, check train.log"
-            )
+            print(f"[{i}/{len(runs)}] {name} — no metrics.csv produced, check train.log")
@@ -50,12 +50,8 @@ PREDICTED_RE = re.compile(
    r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)"
    r"_predicted(?P<local>_local)?\.parquet$"
 )
-SHARD_RE = re.compile(
-    r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)\.(?P<ext>root|parquet)$"
-)
-LEGACY_PREDICTED_RE = re.compile(
-    r"^pbwo4_10000events_hits_predicted(?P<local>_local)?\.parquet$"
-)
+SHARD_RE = re.compile(r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)\.(?P<ext>root|parquet)$")
+LEGACY_PREDICTED_RE = re.compile(r"^pbwo4_10000events_hits_predicted(?P<local>_local)?\.parquet$")
 LEGACY_RE = re.compile(r"^pbwo4_10000events_hits\.(?P<ext>root|parquet)$")


@@ -110,24 +106,9 @@ def plan_moves(src_root: Path) -> tuple[list[tuple[Path, Path]], list[Path]]:
        if m:
            detector, shard, ext = m["detector"], int(m["shard"]), m["ext"]
            if ext == "root":
-                dst = (
-                    src_root
-                    / "raw"
-                    / "steps"
-                    / GEN
-                    / detector
-                    / f"shard-{shard:03d}.root"
-                )
+                dst = src_root / "raw" / "steps" / GEN / detector / f"shard-{shard:03d}.root"
            else:
-                dst = (
-                    src_root
-                    / "processed"
-                    / "steps"
-                    / GEN
-                    / SCHEMA
-                    / detector
-                    / f"shard-{shard:03d}.parquet"
-                )
+                dst = src_root / "processed" / "steps" / GEN / SCHEMA / detector / f"shard-{shard:03d}.parquet"
            moves.append((path, dst))
            continue

@@ -135,19 +116,9 @@ def plan_moves(src_root: Path) -> tuple[list[tuple[Path, Path]], list[Path]]:
        if m:
            ext = m["ext"]
            if ext == "root":
-                dst = (
-                    src_root / "raw" / "hits" / LEGACY_GEN / "pbwo4" / "shard-000.root"
-                )
+                dst = src_root / "raw" / "hits" / LEGACY_GEN / "pbwo4" / "shard-000.root"
            else:
-                dst = (
-                    src_root
-                    / "processed"
-                    / "hits"
-                    / LEGACY_GEN
-                    / LEGACY_SCHEMA
-                    / "pbwo4"
-                    / "shard-000.parquet"
-                )
+                dst = src_root / "processed" / "hits" / LEGACY_GEN / LEGACY_SCHEMA / "pbwo4" / "shard-000.parquet"
            moves.append((path, dst))
            continue

@@ -165,20 +136,9 @@ def plan_manifests(src_root: Path) -> dict[Path, list[str]]:
        for pool, shards in rules.items():
            manifest_path = manifest_dir / f"{pool}{MANIFEST_SUFFIX}"
            for shard in shards:
-                dst = (
-                    src_root
-                    / "processed"
-                    / "steps"
-                    / GEN
-                    / SCHEMA
-                    / detector
-                    / f"shard-{shard:03d}.parquet"
-                )
+                dst = src_root / "processed" / "steps" / GEN / SCHEMA / detector / f"shard-{shard:03d}.parquet"
                manifests[manifest_path].append((shard, dst))
-    return {
-        k: [os.path.relpath(dst, start=k.parent) for _, dst in sorted(v)]
-        for k, v in manifests.items()
-    }
+    return {k: [os.path.relpath(dst, start=k.parent) for _, dst in sorted(v)] for k, v in manifests.items()}


 def run_migration(root: str, execute: bool, copy: bool) -> None:
@@ -13,7 +13,7 @@ real checkpoint) they short-circuit almost instantly and are excluded here —
 see `runtime_estimate.py`'s `_ROUTER_FIXED_S` for how those are handled
 instead.

-Usage: ``uv run python scripts/profile_analysis_costs.py``
+Usage: ``uv run python giant/tools/profile_analysis_costs.py``
 """

 from __future__ import annotations
@@ -26,8 +26,9 @@ import numpy as np
 import polars as pl

 from giant.analysis.catalog import catalog_ids, get_spec
-from giant.analysis.condor import compute_reduced
+from giant.analysis.run import compute_reduced
 from giant.analysis.context import build_context
+from giant.analysis.sources import RolloutSpec

 # Row counts (per side) to benchmark at. Kept in local memory/CPU range so the
 # whole sweep finishes in about a minute; the fit is linear so it extrapolates
@@ -94,9 +95,7 @@ def _make_rollout(n: int, n_events: int, seed: int) -> pl.DataFrame:
            "post_dx": post_dir[:, 0],
            "post_dy": post_dir[:, 1],
            "post_dz": post_dir[:, 2],
-            "edep": np.where(
-                is_synthetic, np.where(reasons == "escaped", 0.0, pre_E), edep
-            ),
+            "edep": np.where(is_synthetic, np.where(reasons == "escaped", 0.0, pre_E), edep),
            "step_length": np.where(is_synthetic, 0.0, step_length),
            "material": rng.choice(_MATERIALS, size=n),
            "layer_id": rng.integers(0, 30, size=n),
@@ -163,17 +162,14 @@ def _make_reference(n: int, n_events: int, seed: int) -> pl.DataFrame:
    )


-def _time(
-    spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path
-) -> float:
+def _time(spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path) -> float:
    t0 = time.perf_counter()
    compute_reduced(
        spec_id,
-        rollout,
+        [{"name": "rollout", "path": str(rollout)}],
        reference,
        shared,
        out,
-        checkpoint=None,
        chunk_index=0,
        n_chunks=1,
    )
@@ -195,7 +191,7 @@ def main() -> None:

            shared = tmp_path / f"shared_{n_side}.json"
            ctx = build_context(
-                rollout,
+                [RolloutSpec(name="rollout", source=rollout)],
                reference,
                n_energy_bins=4,
                n_marginal_bins=50,
@@ -2,11 +2,11 @@

 A single `dwarf convert` call converts a list of files one at a time; this
 module runs up to --jobs conversions concurrently, each as its own `dwarf
-convert` subprocess (invoked via `python -m scripts.dwarf`, so it picks up
+convert` subprocess (invoked via `python -m giant.tools.dwarf`, so it picks up
 the active venv/uv environment automatically).

 Inputs must live under <dataset-root>/raw/<kind>/<gen>/<detector>/<file>.root
-(see scripts/migrate_geant_steps.py) — each is written to the matching
+(see giant/tools/migrate_geant_steps.py) — each is written to the matching
 processed/<kind>/<gen>/<schema>/<detector>/<file>.parquet, where <schema>
 defaults to the highest schemaN already under processed/<kind>/<gen>/ (pass
 --schema to pick a specific one, e.g. one just created by `dwarf bump-schema`).
@@ -22,7 +22,7 @@ import sys
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path

-# Must match scripts/bump_dataset_version.py's GEN_RE / SCHEMA_RE.
+# Must match giant/tools/bump_dataset_version.py's GEN_RE / SCHEMA_RE.
 GEN_RE = re.compile(r"^gen\d+$")
 SCHEMA_RE = re.compile(r"^schema(\d+)$")

@@ -49,9 +49,7 @@ def latest_schema_tag(processed_gen_dir: Path) -> str | None:
    return best_tag


-def resolve_destination(
-    root_file: Path, dataset_root: Path, schema_override: str | None
-) -> Path:
+def resolve_destination(root_file: Path, dataset_root: Path, schema_override: str | None) -> Path:
    """Map raw/<kind>/<gen>/<detector>/<file>.root (relative to *dataset_root*)
    to processed/<kind>/<gen>/<schema>/<detector>/<file>.parquet.

@@ -66,12 +64,7 @@ def resolve_destination(
        raise DestinationError(f"{root_file} is not under dataset root {dataset_root}")

    parts = rel.parts
-    if (
-        len(parts) != 5
-        or parts[0] != "raw"
-        or not GEN_RE.match(parts[2])
-        or not parts[4].endswith(".root")
-    ):
+    if len(parts) != 5 or parts[0] != "raw" or not GEN_RE.match(parts[2]) or not parts[4].endswith(".root"):
        raise DestinationError(
            f"{root_file} does not match raw/<kind>/<gen>/<detector>/<file>.root "
            f"under {dataset_root} (got relative path: {rel})"
@@ -89,7 +82,7 @@ def resolve_destination(
    return processed_gen_dir / schema_tag / detector / f"{shard_stem}.parquet"


-_DWARF_CONVERT_CMD = [sys.executable, "-m", "scripts.dwarf", "convert"]
+_DWARF_CONVERT_CMD = [sys.executable, "-m", "giant.tools.dwarf", "convert"]


 def _convert_one(
@@ -134,7 +127,7 @@ def run_parallel(
    written next to the input .root).

    *cmd_prefix* overrides the subprocess command run per file (defaults to
-    `python -m scripts.dwarf convert`) — used by tests to substitute a fake
+    `python -m giant.tools.dwarf convert`) — used by tests to substitute a fake
    conversion script.

    Returns one (root_file, returncode, stdout, stderr) tuple per file, in
@@ -216,13 +209,7 @@ def run_parallel_job(
            print(f"  {root_file}", file=sys.stderr)
        raise SystemExit(1)

-    total_orphaned = sum(
-        int(m.group(1))
-        for _, _, stdout, _ in results
-        for m in _ORPHAN_RE.finditer(stdout)
-    )
+    total_orphaned = sum(int(m.group(1)) for _, _, stdout, _ in results for m in _ORPHAN_RE.finditer(stdout))
    if total_orphaned:
-        print(
-            f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s)."
-        )
+        print(f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s).")
    print(f"\nAll {len(results)} conversion(s) completed.")
@@ -0,0 +1,98 @@
+"""dwarf warm-cache — precompute `giant train`'s setup-stage sidecar ahead of time.
+
+Thin wrapper around `giant.pipeline.run_setup_stage` so a dataset's vocab
+maps, event-id split index, and normalizer stats can be warmed once — e.g.
+right after `dwarf convert`, or before kicking off a `dwarf hparam-scan`
+sweep — without needing to also start training. See giant/data/setup_cache.py
+for the sidecar itself.
+"""
+
+from pathlib import Path
+
+from giant import config as gconfig
+from giant.pipeline import run_setup_stage
+
+
+def run_warm_setup_cache(
+    data: str,
+    config_path: Path | None = None,
+    val_fraction: float | None = None,
+    seed: int | None = None,
+    particle_conditioning: str | None = None,
+    material_conditioning: str | None = None,
+    router_enabled: bool | None = None,
+    router_type: str | None = None,
+    n_experts: int | None = None,
+    rebuild: bool = False,
+    echo=print,
+) -> None:
+    """Populate (or refresh) the setup cache sidecar for `data`.
+
+    Two mutually exclusive ways to select what to warm for (enforced by the
+    caller, `giant.tools.dwarf.warm_cache` — this function just trusts
+    whichever combination it's given):
+
+    - `config_path`: the same TOML `giant train --config` takes. Every value
+      `run_setup_stage` needs (`train.val_fraction`/`seed`,
+      `conditioning.particle`/`material.type`, both stages' `router`,
+      `stage2_model.particle_type.n_classes`, ...) is read from the one
+      resulting merged `cfg`, so a later `giant train --config <same file>`
+      run resolves to exactly the same cache keys — see gitea #59.
+    - The individual flags below: `val_fraction`/`seed`/
+      `particle_conditioning`/`material_conditioning` select the normalizer
+      cache entry (`giant.data.setup_cache.normalizer_key`) — pass the same
+      values a later `giant train` invocation will use so it hits this
+      warmed entry. The two conditioning axes are independent and may
+      differ. `router_enabled`/`router_type`/`n_experts` only matter for
+      `router_type == "process"` (warms that `n_experts`'s process map); the
+      energy-router quantile summary is always collected regardless, so a
+      later `--router-type energy` run never needs to rescan just to seed
+      centers.
+
+    Any flag left `None` is omitted from the merge, so it falls back to
+    `DEFAULT_CONFIG`'s own value (or the config file's, if `config_path` is
+    given) instead of silently overriding it — see gitea #59.
+    """
+    overrides: dict = {}
+
+    conditioning_overrides: dict = {}
+    if particle_conditioning is not None:
+        conditioning_overrides["particle"] = {"type": particle_conditioning}
+    if material_conditioning is not None:
+        conditioning_overrides["material"] = {"type": material_conditioning}
+    if conditioning_overrides:
+        overrides["conditioning"] = conditioning_overrides
+
+    # This CLI only ever configures one router (matching today's single
+    # --router-type flag), so it's placed on stage1_model; stage2_model's is
+    # left to DEFAULT_CONFIG/the config file rather than forced disabled.
+    router_overrides: dict = {}
+    if router_enabled is not None:
+        router_overrides["enabled"] = router_enabled
+    if router_type is not None:
+        router_overrides["type"] = router_type
+    if n_experts is not None:
+        router_overrides["n_experts"] = n_experts
+    if router_overrides:
+        overrides["stage1_model"] = {"router": router_overrides}
+
+    train_overrides: dict = {}
+    if val_fraction is not None:
+        train_overrides["val_fraction"] = val_fraction
+    if seed is not None:
+        train_overrides["seed"] = seed
+    if train_overrides:
+        overrides["train"] = train_overrides
+
+    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config_path, overrides)
+    gconfig.validate_config(cfg)
+    run_setup_stage(
+        Path(data),
+        val_fraction=cfg["train"]["val_fraction"],
+        seed=cfg["train"]["seed"],
+        cfg=cfg,
+        cache_setup=True,
+        rebuild_setup_cache=rebuild,
+        echo=echo,
+    )
+    echo("setup cache warmed.")
@@ -0,0 +1,31 @@
+"""Training: per-stage trainers, metric collection, checkpointing, the loop.
+
+Split out of the former single-module `giant/train.py`. The public surface is
+`train` (the entry point `giant.pipeline` calls) plus the trainer/spec types
+that tests and tooling construct directly.
+"""
+
+from giant.training.checkpoint import build_checkpoint, init_stages_from_checkpoints, load_checkpoint
+from giant.training.metrics import MetricsCollector, MetricSpec
+from giant.training.loop import train
+from giant.training.trainers import (
+    FlowDDPMStageTrainer,
+    StageSpec,
+    StageTrainer,
+    WGANStageTrainer,
+    build_stage_trainers,
+)
+
+__all__ = [
+    "FlowDDPMStageTrainer",
+    "MetricSpec",
+    "MetricsCollector",
+    "StageSpec",
+    "StageTrainer",
+    "WGANStageTrainer",
+    "build_checkpoint",
+    "build_stage_trainers",
+    "init_stages_from_checkpoints",
+    "load_checkpoint",
+    "train",
+]
@@ -0,0 +1,47 @@
+"""Mixed-precision training support (`train.precision`, gitea #47).
+
+Only `"fp32"` (no autocast) and `"bf16"` are supported — no `"fp16"`/
+`GradScaler`. bf16 needs no gradient scaler and covers every training GPU in
+the fleet (Ampere and newer: A100, L40S, H200, RTX 4070); fp16 would need a
+scaler *and* fixes to two fragile spots that stay correct under bf16 but break
+under fp16's narrower range — `giant.model.routers`' `1e-8` epsilons (below
+fp16's ~6e-8 subnormal floor) and `giant.model.wgan.gradient_penalty`'s
+sum-of-squares gradient norm (overflows fp16 above ~65504). Revisit if a
+pre-Ampere (V100) training target ever shows up.
+"""
+
+import torch
+
+_SUPPORTED_DEVICE_TYPES = ("cuda", "cpu")
+
+
+def resolve_autocast(precision: str, device: torch.device) -> tuple[str, torch.dtype, bool]:
+    """Resolves `train.precision` + a target device into the
+    `(device_type, dtype, enabled)` triple `torch.autocast` takes as kwargs —
+    computed once per `StageTrainer` rather than re-derived every step.
+
+    Raises `ValueError` rather than silently falling back to fp32: a training
+    run that's quietly not using the mixed precision it was configured for is
+    a wasted GPU-week, not a warning.
+    """
+    if precision == "fp32":
+        return device.type, torch.float32, False
+    if precision != "bf16":
+        raise ValueError(f"unknown precision {precision!r}; must be 'fp32' or 'bf16'")
+
+    if device.type == "cuda":
+        if not torch.cuda.is_bf16_supported():
+            cap = torch.cuda.get_device_capability(device)
+            raise ValueError(
+                f"train.precision = 'bf16' but {torch.cuda.get_device_name(device)} "
+                f"(compute capability {cap[0]}.{cap[1]}) has no native bf16 support "
+                "(needs Ampere/sm_80 or newer) — use train.precision = 'fp32' instead"
+            )
+        return "cuda", torch.bfloat16, True
+    if device.type == "cpu":
+        # torch 2.3's CPU autocast supports bf16 unconditionally — this is
+        # also what lets the bf16 training path be tested without a GPU.
+        return "cpu", torch.bfloat16, True
+    raise ValueError(
+        f"train.precision = 'bf16' is not supported on device type {device.type!r} (only {_SUPPORTED_DEVICE_TYPES} are)"
+    )
@@ -0,0 +1,101 @@
+"""Checkpoint assembly and restore.
+
+The on-disk layout is unchanged from v0.2/v0.3.0 and is read by
+`giant/cli.py`, `giant/rollout.py`, `giant/sample.py` and
+`giant/analysis/router_gating.py` — stage 1's weights live under `model`,
+stage 2's under `sec_decoder`, with `_ema`/`critic`/`sec_critic` companions
+and per-stage `optimizer_<stage>` / `optimizer_d_<stage>` / `lr_sched_<stage>`
+entries.
+"""
+
+import torch
+
+from giant.training.trainers import StageTrainer
+
+#: Stage name -> the checkpoint key its weights live under. Historical: stage
+#: 1 predates the two-stage split, so it kept the bare "model" key.
+_STAGE_KEY = {"stage1": "model", "stage2": "sec_decoder"}
+_CRITIC_KEY = {"stage1": "critic", "stage2": "sec_critic"}
+
+
+def build_checkpoint(
+    trainers: dict[str, StageTrainer],
+    epoch: int,
+    global_step: int,
+    best_val_loss: float,
+    extras: dict,
+) -> dict:
+    """`extras` carries the dataset-level sidecars (normalizer, vocab maps,
+    model_config) that `train()` receives as arguments; `None` values are
+    omitted so an absent sidecar leaves no key behind."""
+    ckpt: dict = {
+        "epoch": epoch,
+        "best_val_loss": best_val_loss,
+        "global_step": global_step,
+    }
+    for name, trainer in trainers.items():
+        sd = trainer.state_dict()
+        key = _STAGE_KEY[name]
+        ckpt[key] = sd["model"]
+        if "model_ema" in sd:
+            ckpt[f"{key}_ema"] = sd["model_ema"]
+        if "critic" in sd:
+            ckpt[_CRITIC_KEY[name]] = sd["critic"]
+            ckpt[f"optimizer_d_{name}"] = sd["optimizer_d"]
+        ckpt[f"optimizer_{name}"] = sd["optimizer"]
+        ckpt[f"lr_sched_{name}"] = sd["lr_sched"]
+    ckpt.update({k: v for k, v in extras.items() if v is not None})
+    return ckpt
+
+
+def init_stages_from_checkpoints(trainers: dict[str, StageTrainer]) -> list[str]:
+    """Load each trainer's `spec.init_from` checkpoint (gitea #42) into its
+    model, before training starts — the partial-retrain counterpart to
+    `load_checkpoint`'s full-run `--resume`. Only weights move: unlike
+    `load_checkpoint`, this never touches optimizer/lr_sched/epoch state, so
+    it composes cleanly with `--resume` (call this first; a resume's own
+    `load_checkpoint` then overwrites whatever this loaded with the resumed
+    run's own weights).
+
+    A stage with no `init_from` set (`""`, the default) is left alone. The
+    EMA companion (`<key>_ema`) is loaded too when both the source checkpoint
+    and this trainer have one, so `--weights ema` at inference still sees the
+    source's EMA shadow rather than a copy of its raw weights. Returns one
+    description string per stage actually initialized, for the caller to
+    echo.
+    """
+    loaded = []
+    for name, trainer in trainers.items():
+        init_from = trainer.spec.init_from
+        if not init_from:
+            continue
+        key = _STAGE_KEY[name]
+        ckpt = torch.load(init_from, map_location="cpu", weights_only=False)
+        trainer.model.load_state_dict(ckpt[key])
+        ema_key = f"{key}_ema"
+        if trainer.ema_model is not None and ema_key in ckpt:
+            trainer.ema_model.load_state_dict(ckpt[ema_key])
+        loaded.append(f"{name}: loaded from {init_from}" + (" (frozen)" if trainer.frozen else ""))
+    return loaded
+
+
+def load_checkpoint(trainers: dict[str, StageTrainer], ckpt: dict, lr: float) -> None:
+    """Restore every active stage, then hand `lr`'s authority back to the
+    config — `load_state_dict` would otherwise leave the checkpoint's own
+    base LR in place, silently ignoring `--lr` on resume."""
+    for name, trainer in trainers.items():
+        key = _STAGE_KEY[name]
+        sd = {
+            "model": ckpt[key],
+            "optimizer": ckpt[f"optimizer_{name}"],
+            "lr_sched": ckpt[f"lr_sched_{name}"],
+        }
+        ema_key = f"{key}_ema"
+        if ema_key in ckpt:
+            sd["model_ema"] = ckpt[ema_key]
+        crit_key = _CRITIC_KEY[name]
+        if crit_key in ckpt:
+            sd["critic"] = ckpt[crit_key]
+            sd["optimizer_d"] = ckpt[f"optimizer_d_{name}"]
+        trainer.load_state_dict(sd)
+        trainer.resume_lr(lr)
@@ -0,0 +1,312 @@
+"""The training loop.
+
+`train()` owns the epoch structure and nothing else: the per-stage step is
+`giant.training.trainers`' job, every number reported is
+`giant.training.metrics`' job, and the on-disk checkpoint is
+`giant.training.checkpoint`'s.
+"""
+
+import os
+import signal
+import time
+from pathlib import Path
+from types import FrameType
+from typing import Callable
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from giant import config
+from giant.data.loader import TopNMap
+from giant.data.setup_cache import topnmap_to_json
+from giant.training.checkpoint import build_checkpoint, init_stages_from_checkpoints, load_checkpoint
+from giant.training.metrics import MetricsCollector
+from giant.training.trainers import (
+    FlowDDPMStageTrainer,
+    StageTrainer,
+    build_stage_trainers,
+)
+from giant.validate import validate_marginals
+
+_CATCHABLE_SIGNALS = (signal.SIGINT, signal.SIGTERM)
+
+
+class _GracefulShutdown:
+    """Turns SIGINT/SIGTERM into a flag check instead of an immediate crash.
+
+    A second signal while already shutting down restores the default
+    handler and re-sends the signal, so an unresponsive run can still be
+    force-killed.
+    """
+
+    def __init__(self) -> None:
+        self.requested = False
+        self._previous: dict[
+            int,
+            Callable[[int, FrameType | None], object] | signal.Handlers | int | None,
+        ] = {}
+
+    def __enter__(self) -> "_GracefulShutdown":
+        for sig in _CATCHABLE_SIGNALS:
+            self._previous[sig] = signal.getsignal(sig)
+            signal.signal(sig, self._handle)
+        return self
+
+    def __exit__(self, *exc_info) -> None:
+        for sig, handler in self._previous.items():
+            signal.signal(sig, handler)
+
+    def _handle(self, signum: int, frame) -> None:
+        if self.requested:
+            signal.signal(signum, self._previous[signum])
+            os.kill(os.getpid(), signum)
+            return
+        self.requested = True
+        print(
+            f"\nreceived {signal.Signals(signum).name} — finishing the current "
+            "batch, then saving a checkpoint and exiting (send again to force-quit)"
+        )
+
+
+def _try_validate_marginals(trainer: StageTrainer, val_loader, device, **kwargs):
+    """Runs `validate_marginals` on `trainer`'s sampling model (EMA model if
+    present, else the raw model). `validate_marginals` itself dispatches
+    through `giant.sample.sample_stage1`/`sample_stage2`/`resolve_n_sec`, so
+    this is generator- and one-shot-vs-autoregressive-agnostic."""
+    model = trainer.sampling_model()
+    return validate_marginals(model, val_loader, device=device, **kwargs)
+
+
+def _marginal_kl(trainers: dict[str, StageTrainer], val_loader, device, **kwargs) -> float:
+    """Mean marginal KL over the stage-1 sampling chain, or NaN when stage 1
+    is inactive or `validate_marginals` declined to produce a result."""
+    stage1 = trainers.get("stage1")
+    if stage1 is None:
+        return float("nan")
+    result = _try_validate_marginals(
+        stage1,
+        val_loader,
+        device,
+        sec_decoder=trainers["stage2"].sampling_model() if "stage2" in trainers else None,
+        **kwargs,
+    )
+    if result is None:
+        return float("nan")
+    return float(np.mean(result["kl_divergence"]))
+
+
+def train(
+    cfg: dict,
+    models: dict[str, torch.nn.Module | None],
+    critics: dict[str, torch.nn.Module | None],
+    train_loader: DataLoader,
+    val_loader: DataLoader,
+    device: torch.device,
+    out_dir: str | Path,
+    normalizer_dict: dict | None = None,
+    pdg_map: dict | None = None,
+    mat_map: dict | None = None,
+    proc_map: dict | None = None,
+    pdg_topn_map: TopNMap | None = None,
+    sec_type_topn_map: TopNMap | None = None,
+    mat_topn_map: TopNMap | None = None,
+    model_config: dict | None = None,
+    resume_path: str | Path | None = None,
+    total_train_batches: int = 0,
+    use_wandb: bool = False,
+    wandb_project: str = "giant",
+    wandb_run_name: str = "",
+    wandb_log_every: int = 50,
+) -> None:
+    """Train whichever of stage1/stage2 are active, each through its own
+    `StageTrainer`. `models`/`critics` are the dicts
+    `giant.model.network.build_models`/`build_critics` return — a `None`
+    entry means that stage is `active = false`.
+    """
+    out_dir = Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    t = cfg["train"]
+    epochs = t["epochs"]
+    validate_every = t.get("validate_every", 0)
+    validate_steps = t.get("validate_steps", 10)
+    max_val_batches = t.get("max_val_batches", 0)
+
+    sec_type_class_counts = sec_type_topn_map.class_counts if sec_type_topn_map is not None else None
+    trainers = build_stage_trainers(cfg, models, critics, device, total_train_batches, sec_type_class_counts)
+    if not trainers:
+        raise ValueError("no active stage — stage1_model.active and stage2_model.active are both false")
+    for line in init_stages_from_checkpoints(trainers):
+        print(line)
+    has_adversarial = any(not tr.supports_val_loss for tr in trainers.values())
+
+    checkpoint_extras = {
+        "normalizer": normalizer_dict,
+        "pdg_map": pdg_map,
+        "mat_map": mat_map,
+        "proc_map": proc_map,
+        "pdg_topn_map": topnmap_to_json(pdg_topn_map) if pdg_topn_map is not None else None,
+        "sec_type_topn_map": topnmap_to_json(sec_type_topn_map) if sec_type_topn_map is not None else None,
+        "mat_topn_map": topnmap_to_json(mat_topn_map) if mat_topn_map is not None else None,
+        "model_config": model_config,
+    }
+
+    start_epoch = 1
+    best_val_loss = float("inf")
+    global_step = 0
+    if resume_path is not None:
+        ckpt = torch.load(resume_path, map_location=device, weights_only=False)
+        load_checkpoint(trainers, ckpt, t["lr"])
+        start_epoch = ckpt.get("epoch", 0) + 1
+        best_val_loss = ckpt.get("best_val_loss", float("inf"))
+        global_step = ckpt.get("global_step", 0)
+        if start_epoch > epochs:
+            print(f"checkpoint already completed epoch {start_epoch - 1} (>= --epochs {epochs}) — nothing to train")
+            return
+
+    collector = MetricsCollector.create(
+        trainers,
+        out_dir,
+        cfg,
+        model_config,
+        resume=resume_path is not None,
+        use_wandb=use_wandb,
+        wandb_project=wandb_project,
+        wandb_run_name=wandb_run_name,
+        wandb_log_every=wandb_log_every,
+    )
+
+    epoch_w = len(str(epochs))
+    last_completed_epoch = start_epoch - 1
+    with _GracefulShutdown() as shutdown:
+        for epoch in range(start_epoch, epochs + 1):
+            epoch_start = time.monotonic()
+            if device.type == "cuda":
+                torch.cuda.reset_peak_memory_stats(device)
+            collector.start_epoch(epoch)
+            # Epoch-aware RNG: same noise (and, below, same batch order) for
+            # epoch k whether the run is one process or a chain of per-epoch
+            # jobs. See giant.config.epoch_seed.
+            config.seed_everything(config.epoch_seed(t["seed"], epoch))
+            # Epoch-aware shuffle stream (see StreamingStepsDataset.set_epoch):
+            # keeps epoch k's batch order identical whether it runs here or as
+            # its own resumed per-epoch job in a b2luigi workflow.
+            # (tests hand `train` a plain list of batches, which has neither)
+            set_epoch = getattr(getattr(train_loader, "dataset", None), "set_epoch", None)
+            if callable(set_epoch):
+                set_epoch(epoch)
+            for trainer in trainers.values():
+                trainer.train_mode()
+
+            bar = tqdm(
+                train_loader,
+                desc=f"  epoch {epoch:{epoch_w}d}/{epochs}",
+                total=total_train_batches or None,
+                leave=False,
+                unit="batch",
+                dynamic_ncols=True,
+            )
+            for batch in bar:
+                B = batch[0].size(0)
+                collector.add_train_batch(
+                    {name: trainer.step(batch, device, global_step) for name, trainer in trainers.items()},
+                    B,
+                )
+                bar.set_postfix_str(collector.postfix(), refresh=False)
+                global_step += 1
+                collector.log_batch(global_step, batch, device)
+                if shutdown.requested:
+                    break
+            bar.close()
+
+            if shutdown.requested:
+                ckpt = build_checkpoint(trainers, epoch - 1, global_step, best_val_loss, checkpoint_extras)
+                torch.save(ckpt, out_dir / "last.pt")
+                last_completed_epoch = epoch - 1
+                print(
+                    f"saved in-progress weights from partway through epoch "
+                    f"{epoch} to {out_dir / 'last.pt'} "
+                    f"(resume will restart epoch {epoch})"
+                )
+                break
+
+            for trainer in trainers.values():
+                trainer.eval_mode()
+
+            # --- per-stage validation ---
+            scored = {name: tr for name, tr in trainers.items() if tr.supports_val_loss}
+            if scored:
+                with torch.no_grad():
+                    for val_batch_idx, batch in enumerate(val_loader):
+                        if max_val_batches > 0 and val_batch_idx >= max_val_batches:
+                            break
+                        B = batch[0].size(0)
+                        collector.add_val_batch(
+                            {name: tr.val_loss(batch, device) for name, tr in scored.items()},
+                            B,
+                        )
+                        collector.observe_routers(batch[0].to(device), batch[1].to(device), B)
+
+            # An adversarial stage has no averageable validation loss, so it
+            # needs the marginal-KL signal every epoch to pick a best
+            # checkpoint at all; a purely non-adversarial run only pays for
+            # it every `validate_every` epochs.
+            marginal_kl = float("nan")
+            if has_adversarial:
+                marginal_kl = _marginal_kl(trainers, val_loader, device)
+            elif validate_every > 0 and epoch % validate_every == 0:
+                stage1 = trainers.get("stage1")
+                ddpm_steps = 1000
+                if isinstance(stage1, FlowDDPMStageTrainer) and stage1.ddpm_schedule is not None:
+                    ddpm_steps = stage1.ddpm_schedule.T
+                marginal_kl = _marginal_kl(
+                    trainers,
+                    val_loader,
+                    device,
+                    steps=validate_steps,
+                    ddpm_steps=ddpm_steps,
+                )
+
+            val_loss = sum(
+                trainer.val_objective(
+                    collector.train_means(name),
+                    collector.val_means(name),
+                    marginal_kl,
+                )
+                for name, trainer in trainers.items()
+            )
+
+            epoch_time = time.monotonic() - epoch_start
+            is_best = val_loss < best_val_loss
+            collector.set("val/loss", val_loss)
+            collector.set("val/marginal_kl", marginal_kl)
+            collector.set(
+                "gpu_mem_mb",
+                torch.cuda.max_memory_allocated(device) / (1024 * 1024) if device.type == "cuda" else 0.0,
+            )
+            collector.set("samples_per_sec", collector.train_samples / max(epoch_time, 1e-8))
+            collector.set("is_best", int(is_best))
+            collector.set("epoch_time_s", epoch_time)
+
+            print(collector.summary_line(val_loss, epoch_time, is_best))
+            collector.write_epoch(global_step)
+
+            ckpt = build_checkpoint(trainers, epoch, global_step, best_val_loss, checkpoint_extras)
+            if is_best:
+                best_val_loss = val_loss
+                ckpt["best_val_loss"] = best_val_loss
+                torch.save(ckpt, out_dir / "best.pt")
+            torch.save(ckpt, out_dir / "last.pt")
+            last_completed_epoch = epoch
+
+            if shutdown.requested:
+                break
+
+    collector.close()
+
+    if shutdown.requested:
+        print(
+            f"stopped after epoch {last_completed_epoch} due to shutdown signal — "
+            f"resume with --resume {out_dir / 'last.pt'}"
+        )
@@ -0,0 +1,384 @@
+"""Per-epoch metric accumulation, `metrics.csv`, and W&B logging.
+
+Every scalar a training run reports is declared exactly once, as a
+`MetricSpec` on the `StageTrainer` that computes it (see
+`giant.training.trainers`). `MetricsCollector` derives the CSV/W&B column set
+from those declarations, so adding a metric means adding one line next to the
+code that produces it — there is no second list to keep in sync.
+
+Column naming is uniform: `<stage>/train/<key>`, `<stage>/val/<key>`,
+`<stage>/<key>` for point-in-time values (`lr`, `critic_lr`),
+`<stage>/router/<key>` for routing diagnostics, and an unprefixed run-level
+tail (`val/loss`, `grad_norm`, `epoch_time_s`, ...). W&B groups panels on
+`/`, so the same names read well there.
+"""
+
+import csv
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+
+_ROUTER_KEYS = ("entropy", "util_min", "util_max", "util_std")
+
+# Written after every stage's columns, by `MetricsCollector` itself rather
+# than by any one trainer — these describe the run, not a stage.
+_RUN_COLUMNS = (
+    "val/loss",
+    "val/marginal_kl",
+    "grad_norm",
+    "gpu_mem_mb",
+    "samples_per_sec",
+    "is_best",
+    "epoch_time_s",
+)
+
+# tqdm/W&B batch-granularity smoothing, matching v0.2/v0.3.0's inline EMA.
+_EMA_ALPHA = 0.05
+
+
+@dataclass(frozen=True)
+class MetricSpec:
+    """One scalar a trainer emits per batch, and how it is reported.
+
+    `key` indexes the dict `StageTrainer.step()` / `.val_loss()` returns;
+    `column` is the CSV/W&B column suffix, joined to the stage name with
+    "/". `reduce` is either "mean" (batch-size-weighted average over the
+    epoch) or "last" (the most recent value — for point-in-time quantities
+    like the learning rate, which is a schedule readout, not a statistic).
+    """
+
+    key: str
+    column: str
+    reduce: str = "mean"
+
+
+def train_metric(key: str, column: str | None = None) -> MetricSpec:
+    return MetricSpec(key, column or f"train/{key}")
+
+
+def val_metric(key: str, column: str | None = None) -> MetricSpec:
+    return MetricSpec(key, column or f"val/{key}")
+
+
+def stage_metric(key: str, column: str | None = None) -> MetricSpec:
+    """A point-in-time stage-level readout (`lr`, `critic_lr`) — reported
+    unprefixed by split, as `<stage>/<key>`."""
+    return MetricSpec(key, column or key, reduce="last")
+
+
+def _wandb_run_config(cfg: dict, model_config: dict | None, param_counts: dict) -> dict:
+    return {
+        "train": cfg["train"],
+        "conditioning": cfg["conditioning"],
+        "stage1_model": cfg["stage1_model"],
+        "stage2_model": cfg["stage2_model"],
+        "model_config": model_config or {},
+        "param_counts": param_counts,
+    }
+
+
+class _Accumulator:
+    """Batch-size-weighted sums for one stage and one split."""
+
+    def __init__(self) -> None:
+        self.sums: dict[str, float] = {}
+        self.n = 0
+        self.last: dict[str, float] = {}
+
+    def add(self, stats: dict, keys: set[str], batch_size: int) -> None:
+        for key in keys:
+            if key in stats:
+                self.sums[key] = self.sums.get(key, 0.0) + stats[key] * batch_size
+        self.last.update(stats)
+        self.n += batch_size
+
+    def mean(self, key: str) -> float:
+        return self.sums.get(key, 0.0) / max(self.n, 1)
+
+    def means(self) -> dict[str, float]:
+        return {key: self.mean(key) for key in self.sums}
+
+    def reset(self) -> None:
+        self.sums.clear()
+        self.last.clear()
+        self.n = 0
+
+
+class _RouterAccumulator:
+    """Gate-diagnostic sums for one routed stage."""
+
+    def __init__(self, n_experts: int) -> None:
+        self.n_experts = n_experts
+        self.entropy = 0.0
+        self.importance: torch.Tensor | None = None
+        self.n = 0
+
+    def add(self, entropy: torch.Tensor, importance: torch.Tensor, n: int) -> None:
+        self.entropy += entropy.item() * n
+        self.importance = importance.clone() if self.importance is None else self.importance + importance
+        self.n += n
+
+    def stats(self) -> dict[str, float]:
+        if self.importance is None or self.n == 0:
+            return dict.fromkeys(_ROUTER_KEYS, 0.0)
+        util = self.importance / self.importance.sum().clamp_min(1e-8)
+        return {
+            "entropy": self.entropy / self.n,
+            "util_min": util.min().item(),
+            "util_max": util.max().item(),
+            "util_std": util.std().item() if self.n_experts > 1 else 0.0,
+        }
+
+    def reset(self) -> None:
+        self.entropy = 0.0
+        self.importance = None
+        self.n = 0
+
+
+class MetricsCollector:
+    """Owns every number a training run reports.
+
+    Accumulates per-batch stats from each stage, writes one `metrics.csv` row
+    per epoch, mirrors it to W&B, and formats the tqdm postfix and the epoch
+    summary line — so `giant.training.loop.train` never carries a running
+    sum, a column name, or a W&B call of its own.
+    """
+
+    def __init__(
+        self,
+        trainers: dict,
+        out_dir: Path,
+        *,
+        epochs: int,
+        resume: bool = False,
+        wandb_run=None,
+        wandb_log_every: int = 50,
+    ) -> None:
+        self.trainers = trainers
+        self.epochs = epochs
+        self.wandb_run = wandb_run
+        self.wandb_log_every = wandb_log_every
+        self.epoch_width = len(str(epochs))
+
+        self._train = {name: _Accumulator() for name in trainers}
+        self._val = {name: _Accumulator() for name in trainers}
+        self._routers = {
+            name: _RouterAccumulator(tr.router.n_experts) for name, tr in trainers.items() if tr.router is not None
+        }
+        # Only "mean" specs need summing; "last" specs are read straight off
+        # the accumulator's most recent stats dict. "grad_norm" is always
+        # summed — it feeds the run-level `grad_norm` column whether or not
+        # a trainer reports it per stage.
+        self._train_keys = {
+            name: {spec.key for spec in tr.train_metrics if spec.reduce == "mean"} | {"grad_norm"}
+            for name, tr in trainers.items()
+        }
+        self._val_keys = {
+            name: {spec.key for spec in tr.val_metrics if spec.reduce == "mean"} for name, tr in trainers.items()
+        }
+        self._run_values: dict[str, float] = {}
+        self._epoch = 0
+        self._ema_loss = 0.0
+        self._ema_grad_norm = 0.0
+        self._ema_seeded = False
+        self._batch_loss = 0.0
+        self._batch_grad_norm = 0.0
+
+        self.fieldnames = self._build_fieldnames()
+        metrics_path = out_dir / "metrics.csv"
+        append = resume and metrics_path.exists()
+        self._file = open(metrics_path, "a" if append else "w", newline="")
+        self._writer = csv.DictWriter(self._file, fieldnames=self.fieldnames)
+        if not append:
+            self._writer.writeheader()
+
+    # --- construction ---------------------------------------------------
+
+    @classmethod
+    def create(
+        cls,
+        trainers: dict,
+        out_dir: Path,
+        cfg: dict,
+        model_config: dict | None,
+        *,
+        resume: bool = False,
+        use_wandb: bool = False,
+        wandb_project: str = "giant",
+        wandb_run_name: str = "",
+        wandb_log_every: int = 50,
+    ) -> "MetricsCollector":
+        """Build the collector, starting a W&B run first when enabled."""
+        wandb_run = None
+        if use_wandb:
+            try:
+                import wandb
+            except ImportError as exc:
+                raise RuntimeError(
+                    "train.wandb = true (--wandb) requires the 'wandb' package — install it via `uv sync --extra wandb`"
+                ) from exc
+            param_counts = {name: sum(p.numel() for p in tr.model.parameters()) for name, tr in trainers.items()}
+            param_counts["total"] = sum(param_counts.values())
+            wandb_run = wandb.init(
+                project=wandb_project,
+                name=wandb_run_name or out_dir.name,
+                id=out_dir.name,
+                resume="allow",
+                config=_wandb_run_config(cfg, model_config, param_counts),
+            )
+        return cls(
+            trainers,
+            out_dir,
+            epochs=cfg["train"]["epochs"],
+            resume=resume,
+            wandb_run=wandb_run,
+            wandb_log_every=wandb_log_every,
+        )
+
+    def _build_fieldnames(self) -> list[str]:
+        fields = ["epoch"]
+        for name, trainer in self.trainers.items():
+            for spec in trainer.train_metrics:
+                fields.append(f"{name}/{spec.column}")
+            for spec in trainer.val_metrics:
+                fields.append(f"{name}/{spec.column}")
+            if trainer.router is not None:
+                fields += [f"{name}/router/{key}" for key in _ROUTER_KEYS]
+            for spec in trainer.stage_metrics:
+                fields.append(f"{name}/{spec.column}")
+        fields += list(_RUN_COLUMNS)
+        return fields
+
+    def close(self) -> None:
+        self._file.close()
+        if self.wandb_run is not None:
+            self.wandb_run.finish()
+
+    # --- per-batch ------------------------------------------------------
+
+    def start_epoch(self, epoch: int) -> None:
+        self._epoch = epoch
+        for acc in self._train.values():
+            acc.reset()
+        for acc in self._val.values():
+            acc.reset()
+        for acc in self._routers.values():
+            acc.reset()
+        self._run_values.clear()
+        self._ema_seeded = False
+
+    def add_train_batch(self, stats: dict[str, dict], batch_size: int) -> None:
+        """`stats` maps stage name -> the dict that stage's `step()` returned."""
+        self._batch_loss = 0.0
+        self._batch_grad_norm = 0.0
+        for name, stage_stats in stats.items():
+            self._train[name].add(stage_stats, self._train_keys[name], batch_size)
+            self._batch_loss += self.trainers[name].batch_loss(stage_stats)
+            self._batch_grad_norm += stage_stats.get("grad_norm", 0.0)
+        if self._ema_seeded:
+            self._ema_loss += _EMA_ALPHA * (self._batch_loss - self._ema_loss)
+            self._ema_grad_norm += _EMA_ALPHA * (self._batch_grad_norm - self._ema_grad_norm)
+        else:
+            self._ema_loss = self._batch_loss
+            self._ema_grad_norm = self._batch_grad_norm
+            self._ema_seeded = True
+
+    def add_val_batch(self, stats: dict[str, dict], batch_size: int) -> None:
+        for name, stage_stats in stats.items():
+            self._val[name].add(stage_stats, self._val_keys[name], batch_size)
+
+    @torch.no_grad()
+    def observe_routers(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, batch_size: int) -> None:
+        """Record gate diagnostics for every routed stage on this batch.
+
+        Called from the validation pass only (as in v0.2/v0.3.0), so a stage
+        whose trainer has no validation pass — i.e. WGAN — reports zeros.
+        """
+        for name, acc in self._routers.items():
+            router = self.trainers[name].router
+            entropy, importance = router.gate_stats(cond_cont, cond_cat)
+            acc.add(entropy, importance, batch_size)
+
+    def postfix(self) -> str:
+        """tqdm postfix for the training bar."""
+        return f"loss={self._ema_loss:.4f} gnorm={self._ema_grad_norm:.3f}"
+
+    def log_batch(self, global_step: int, batch: tuple, device: torch.device) -> None:
+        """Batch-granularity W&B log, throttled to every `wandb_log_every`
+        optimizer steps (a single epoch can be tens of thousands). `batch` is
+        the raw training batch, needed only to re-derive routing entropy for
+        routed stages — it is never moved to `device` otherwise."""
+        if self.wandb_run is None or self.wandb_log_every <= 0:
+            return
+        if global_step % self.wandb_log_every != 0:
+            return
+        payload = {
+            "batch/epoch": self._epoch,
+            "batch/loss": self._batch_loss,
+            "batch/loss_ema": self._ema_loss,
+            "batch/grad_norm": self._batch_grad_norm,
+        }
+        for name, trainer in self.trainers.items():
+            payload[f"batch/{name}/lr"] = trainer.optimizer.param_groups[0]["lr"]
+            if trainer.router is not None:
+                with torch.no_grad():
+                    entropy, _ = trainer.router.gate_stats(batch[0].to(device), batch[1].to(device))
+                payload[f"batch/{name}/router/entropy"] = entropy.item()
+        self.wandb_run.log(payload, step=global_step)
+
+    # --- per-epoch ------------------------------------------------------
+
+    def train_means(self, stage: str) -> dict[str, float]:
+        return self._train[stage].means()
+
+    def val_means(self, stage: str) -> dict[str, float]:
+        return self._val[stage].means()
+
+    @property
+    def train_samples(self) -> int:
+        """Samples seen this epoch — identical across stages (every stage
+        steps on every batch), so any one accumulator's count will do."""
+        return max((acc.n for acc in self._train.values()), default=0)
+
+    def set(self, column: str, value: float) -> None:
+        """Record a run-level value for this epoch's row (`val/loss`,
+        `gpu_mem_mb`, ...). Must name a column in `_RUN_COLUMNS`."""
+        if column not in _RUN_COLUMNS:
+            raise KeyError(f"{column!r} is not a run-level metrics column")
+        self._run_values[column] = value
+
+    def summary_line(self, val_loss: float, epoch_time: float, is_best: bool) -> str:
+        bits = [trainer.summary(self.train_means(name)) for name, trainer in self.trainers.items()]
+        marker = "  [best]" if is_best else ""
+        return (
+            f"epoch {self._epoch:{self.epoch_width}d}/{self.epochs}  "
+            + "  ".join(bits)
+            + f"  val {val_loss:.4f}  {epoch_time:.1f}s{marker}"
+        )
+
+    def write_epoch(self, global_step: int) -> None:
+        """Assemble, write, and flush this epoch's row; mirror it to W&B."""
+        row: dict = {"epoch": self._epoch}
+        grad_norm_total = 0.0
+        for name, trainer in self.trainers.items():
+            train_acc, val_acc = self._train[name], self._val[name]
+            for spec in trainer.train_metrics:
+                row[f"{name}/{spec.column}"] = train_acc.mean(spec.key)
+            for spec in trainer.val_metrics:
+                row[f"{name}/{spec.column}"] = val_acc.mean(spec.key)
+            if trainer.router is not None:
+                for key, value in self._routers[name].stats().items():
+                    row[f"{name}/router/{key}"] = value
+            for spec in trainer.stage_metrics:
+                row[f"{name}/{spec.column}"] = train_acc.last.get(spec.key, 0.0)
+            grad_norm_total += train_acc.mean("grad_norm")
+
+        for column in _RUN_COLUMNS:
+            row[column] = self._run_values.get(column, float("nan"))
+        row["grad_norm"] = grad_norm_total
+
+        self._writer.writerow(row)
+        self._file.flush()
+        if self.wandb_run is not None:
+            self.wandb_run.log(row, step=global_step)
@@ -0,0 +1,356 @@
+"""Training-progress plots from `<run_dir>/metrics.csv` (gitea #75).
+
+`MetricsCollector` (`giant.training.metrics`) writes one row per epoch with a
+column set that varies by run — flow/ddpm vs wgan, routed vs not (see the
+`MetricSpec` declarations in `giant.training.trainers`). This module reads
+that header dynamically rather than hardcoding a column list, buckets columns
+by the fixed naming convention `MetricsCollector` itself documents
+(`<stage>/train/<key>`, `<stage>/val/<key>`, `<stage>/router/<key>`,
+`<stage>/<key>` for point-in-time values, and an unprefixed run-level tail —
+see `giant.training.metrics`'s module docstring), and renders one PDF per
+applicable figure with the same `plotstyle` conventions
+`giant.analysis.render` uses, for visual consistency with the
+rollout-vs-reference plots.
+
+Unlike `giant.analysis`, there is no reduce/chunk/condor split here — the CSV
+is tiny and this always runs as one local pass — but the CLI entry point
+still lives under `giant analyze` (`analyze metrics`) as the shared home for
+plotstyle-rendered diagnostics, and shares its `analysis_runs/` output
+convention (see `derive_metrics_dir`) so training-progress plots don't get
+written into the training run directory itself.
+"""
+
+from __future__ import annotations
+
+import csv
+import math
+from dataclasses import dataclass
+from pathlib import Path
+
+# Stage names are always exactly these two — hardcoded in
+# `giant.training.trainers.build_stage_trainers` — so a column belongs to a
+# stage iff it's prefixed by one of these, and everything else (bar `epoch`)
+# is run-level. This is what makes dynamic header parsing tractable without
+# needing to know the per-run metric keys themselves.
+_STAGE_NAMES = ("stage1", "stage2")
+
+_ACC_KEYS = {"nsec_acc", "stop_acc", "type_acc"}
+_WGAN_BALANCE_KEYS = {"d_loss", "g_loss", "wasserstein", "gp_loss"}
+_ROUTER_KEYS = ("entropy", "util_min", "util_max", "util_std")
+
+
+@dataclass
+class MetricsTable:
+    """`<run_dir>/metrics.csv`, parsed with no hardcoded column list."""
+
+    epochs: list[int]
+    columns: dict[str, list[float]]
+
+    @classmethod
+    def load(cls, path: str | Path) -> "MetricsTable":
+        with open(path, newline="") as f:
+            rows = list(csv.DictReader(f))
+        epochs = [int(float(r["epoch"])) for r in rows]
+        fieldnames = rows[0].keys() if rows else []
+        columns = {name: [float(r[name]) for r in rows] for name in fieldnames if name != "epoch"}
+        return cls(epochs=epochs, columns=columns)
+
+    def best_epochs(self) -> list[int]:
+        is_best = self.columns.get("is_best")
+        if not is_best:
+            return []
+        return [epoch for epoch, flag in zip(self.epochs, is_best) if flag]
+
+
+# --- column classification --------------------------------------------------
+
+
+def _stages(columns: dict) -> list[str]:
+    return [s for s in _STAGE_NAMES if any(name.startswith(f"{s}/") for name in columns)]
+
+
+def _split(columns: dict, stage: str, split: str) -> dict[str, str]:
+    prefix = f"{stage}/{split}/"
+    return {name[len(prefix) :]: name for name in columns if name.startswith(prefix)}
+
+
+def _point_in_time(columns: dict, stage: str) -> dict[str, str]:
+    prefix = f"{stage}/"
+    out = {}
+    for name in columns:
+        if not name.startswith(prefix):
+            continue
+        rest = name[len(prefix) :]
+        head = rest.split("/", 1)[0]
+        if head not in ("train", "val", "router"):
+            out[rest] = name
+    return out
+
+
+def _router(columns: dict, stage: str) -> dict[str, str]:
+    prefix = f"{stage}/router/"
+    return {name[len(prefix) :]: name for name in columns if name.startswith(prefix)}
+
+
+def _run_level(columns: dict) -> dict[str, str]:
+    known_prefixes = tuple(f"{s}/" for s in _STAGE_NAMES)
+    return {name: name for name in columns if not name.startswith(known_prefixes)}
+
+
+def _loss_keys(train: dict[str, str], val: dict[str, str]) -> list[str]:
+    keys = {k for k in train if k not in _ACC_KEYS and k not in _WGAN_BALANCE_KEYS and k != "grad_norm"}
+    keys |= {k for k in val if k not in _ACC_KEYS and k not in _WGAN_BALANCE_KEYS and k != "grad_norm"}
+    return sorted(keys)
+
+
+# --- output location ---------------------------------------------------------
+
+
+def derive_metrics_dir(
+    run_dir: str | Path,
+    out_dir: str | Path | None = None,
+    default_base: str | Path | None = None,
+) -> Path:
+    """Plots output directory.
+
+    Precedence: an explicit `out_dir` always wins. Otherwise
+    `default_base / f"metrics_{run_dir.name}"` (the CLI passes the repo's
+    gitignored `analysis_runs/`, matching `giant.analysis.run.derive_run_dir`'s
+    convention) — training-progress plots live alongside rollout-vs-reference
+    analysis runs, not inside the training run directory itself.
+    """
+    if out_dir is not None:
+        return Path(out_dir)
+    base = Path(default_base) if default_base is not None else Path.cwd() / "analysis_runs"
+    return base / f"metrics_{Path(run_dir).name}"
+
+
+# --- figures ------------------------------------------------------------------
+
+
+def _mark_best(ax, table: MetricsTable) -> None:
+    for epoch in table.best_epochs():
+        ax.axvline(epoch, color="grey", linestyle="--", linewidth=0.8, alpha=0.7)
+
+
+def _overview_figure(table: MetricsTable):
+    import plotstyle as ps
+
+    run_level = _run_level(table.columns)
+    if "val/loss" not in run_level:
+        return None
+    fig, ax = ps.new_figure("thesis-single", title="training overview")
+    ax.plot(table.epochs, table.columns["val/loss"], label="val/loss")
+    if "val/marginal_kl" in run_level:
+        kl = table.columns["val/marginal_kl"]
+        if any(math.isfinite(v) for v in kl):
+            ax.plot(table.epochs, kl, label="val/marginal_kl")
+    _mark_best(ax, table)
+    best = table.best_epochs()
+    if best:
+        idx = table.epochs.index(best[-1])
+        ax.annotate(
+            f"best: epoch {best[-1]}\nval/loss={table.columns['val/loss'][idx]:.4g}",
+            xy=(best[-1], table.columns["val/loss"][idx]),
+            xytext=(0.98, 0.95),
+            textcoords="axes fraction",
+            ha="right",
+            va="top",
+            fontsize=8,
+        )
+    ax.set_xlabel("epoch")
+    ax.set_ylabel("loss")
+    ps.style_legend(ax, title="series")
+    return fig
+
+
+def _loss_figure(table: MetricsTable, stage: str):
+    import plotstyle as ps
+
+    train = _split(table.columns, stage, "train")
+    val = _split(table.columns, stage, "val")
+    keys = _loss_keys(train, val)
+    if not keys:
+        return None
+    n = len(keys)
+    ncols = min(3, n)
+    nrows = (n + ncols - 1) // ncols
+    fig, axes = ps.new_figure(
+        "slide-16x9",
+        title=f"{stage} loss",
+        nrows=nrows,
+        ncols=ncols,
+        squeeze=False,
+    )
+    flat = axes.ravel()
+    for ax, key in zip(flat, keys):
+        if key in train:
+            ax.plot(table.epochs, table.columns[train[key]], label="train")
+        if key in val:
+            ax.plot(table.epochs, table.columns[val[key]], label="val")
+        ax.set_yscale("log")
+        ax.set_title(key, fontsize=8)
+        ax.set_xlabel("epoch")
+    for j in range(n, len(flat)):
+        flat[j].set_visible(False)
+    ps.style_legend(flat[0], title="series")
+    return fig
+
+
+def _lr_figure(table: MetricsTable):
+    import plotstyle as ps
+
+    series: dict[str, str] = {}
+    for stage in _stages(table.columns):
+        for key, col in _point_in_time(table.columns, stage).items():
+            series[f"{stage}/{key}"] = col
+    if not series:
+        return None
+    fig, ax = ps.new_figure("thesis-single", title="learning rate schedule")
+    for label, col in series.items():
+        ax.plot(table.epochs, table.columns[col], label=label)
+    ax.set_xlabel("epoch")
+    ax.set_ylabel("learning rate")
+    ps.style_legend(ax, title="series")
+    return fig
+
+
+def _accuracy_figure(table: MetricsTable, stage: str):
+    import plotstyle as ps
+
+    train = _split(table.columns, stage, "train")
+    val = _split(table.columns, stage, "val")
+    keys = sorted((set(train) | set(val)) & _ACC_KEYS)
+    if not keys:
+        return None
+    n = len(keys)
+    fig, axes = ps.new_figure("slide-16x9", title=f"{stage} accuracy", nrows=1, ncols=n, squeeze=False)
+    flat = axes.ravel()
+    for ax, key in zip(flat, keys):
+        if key in train:
+            ax.plot(table.epochs, table.columns[train[key]], label="train")
+        if key in val:
+            ax.plot(table.epochs, table.columns[val[key]], label="val")
+        ax.set_title(key, fontsize=8)
+        ax.set_xlabel("epoch")
+        ax.set_ylim(0, 1)
+    ps.style_legend(flat[0], title="series")
+    return fig
+
+
+def _grad_norm_figure(table: MetricsTable):
+    import plotstyle as ps
+
+    run_level = _run_level(table.columns)
+    if "grad_norm" not in run_level:
+        return None
+    fig, ax = ps.new_figure("thesis-single", title="gradient norm")
+    ax.plot(table.epochs, table.columns["grad_norm"], label="grad_norm")
+    for stage in _stages(table.columns):
+        train = _split(table.columns, stage, "train")
+        for key in ("grad_norm_d", "grad_norm_g", "grad_norm_type_slice", "grad_norm_cont_slice"):
+            if key in train:
+                ax.plot(table.epochs, table.columns[train[key]], label=f"{stage}/{key}")
+    ax.set_yscale("log")
+    ax.set_xlabel("epoch")
+    ax.set_ylabel("grad norm")
+    ps.style_legend(ax, title="series")
+    return fig
+
+
+def _router_figure(table: MetricsTable, stage: str):
+    import plotstyle as ps
+
+    router = _router(table.columns, stage)
+    if "entropy" not in router:
+        return None
+    fig, ax = ps.new_figure("thesis-single", title=f"{stage} router health")
+    ax.plot(table.epochs, table.columns[router["entropy"]], label="entropy", color="black")
+    ax.set_xlabel("epoch")
+    ax.set_ylabel("entropy [bits]")
+    ax2 = ax.twinx()
+    for key in ("util_min", "util_max", "util_std"):
+        if key in router:
+            ax2.plot(table.epochs, table.columns[router[key]], label=key, linestyle="--")
+    ax2.set_ylabel("expert utilization")
+    ax2.set_ylim(0, 1)
+    lines1, labels1 = ax.get_legend_handles_labels()
+    lines2, labels2 = ax2.get_legend_handles_labels()
+    ax.legend(lines1 + lines2, labels1 + labels2, loc="upper right", frameon=False, fontsize=7)
+    return fig
+
+
+def _wgan_balance_figure(table: MetricsTable, stage: str):
+    import plotstyle as ps
+
+    train = _split(table.columns, stage, "train")
+    keys = [k for k in _WGAN_BALANCE_KEYS if k in train]
+    if not keys:
+        return None
+    fig, ax = ps.new_figure("thesis-single", title=f"{stage} WGAN critic/generator balance")
+    for key in sorted(keys):
+        ax.plot(table.epochs, table.columns[train[key]], label=key)
+    ax.set_xlabel("epoch")
+    ax.set_ylabel("value")
+    ps.style_legend(ax, title="series")
+    return fig
+
+
+def _throughput_figure(table: MetricsTable):
+    import plotstyle as ps
+
+    run_level = _run_level(table.columns)
+    keys = [k for k in ("samples_per_sec", "gpu_mem_mb", "epoch_time_s") if k in run_level]
+    if not keys:
+        return None
+    fig, axes = ps.new_figure("slide-16x9", title="throughput / resources", nrows=1, ncols=len(keys), squeeze=False)
+    flat = axes.ravel()
+    for ax, key in zip(flat, keys):
+        ax.plot(table.epochs, table.columns[key])
+        _mark_best(ax, table)
+        ax.set_title(key, fontsize=8)
+        ax.set_xlabel("epoch")
+    return fig
+
+
+# --- entry point ---------------------------------------------------------
+
+
+def render_metrics(
+    run_dir: str | Path,
+    out_dir: str | Path | None = None,
+    default_base: str | Path | None = None,
+) -> list[Path]:
+    """`<run_dir>/metrics.csv` -> `<plots dir>/<name>.pdf`.
+
+    See `derive_metrics_dir` for how the plots directory is resolved.
+    """
+    import matplotlib.pyplot as plt
+    import plotstyle as ps
+
+    ps.use()
+    table = MetricsTable.load(Path(run_dir) / "metrics.csv")
+    plots_dir = derive_metrics_dir(run_dir, out_dir, default_base)
+    plots_dir.mkdir(parents=True, exist_ok=True)
+
+    figures = [("overview", _overview_figure(table))]
+    for stage in _stages(table.columns):
+        figures.append((f"{stage}_loss", _loss_figure(table, stage)))
+    figures.append(("lr", _lr_figure(table)))
+    for stage in _stages(table.columns):
+        figures.append((f"{stage}_accuracy", _accuracy_figure(table, stage)))
+    figures.append(("grad_norm", _grad_norm_figure(table)))
+    for stage in _stages(table.columns):
+        figures.append((f"{stage}_router", _router_figure(table, stage)))
+        figures.append((f"{stage}_wgan_balance", _wgan_balance_figure(table, stage)))
+    figures.append(("throughput", _throughput_figure(table)))
+
+    paths: list[Path] = []
+    for name, fig in figures:
+        if fig is None:
+            continue
+        path = plots_dir / name
+        ps.savefig(fig, str(path), formats=("pdf",))
+        plt.close(fig)
+        paths.append(path.with_suffix(".pdf"))
+    return paths
@@ -0,0 +1,366 @@
+"""Ground-truth tensor assembly for stage-2 training.
+
+Pure functions, no optimizer/model state: they turn a batch's ground-truth
+secondary tensors into the per-token targets and autoregressive conditioning
+inputs `giant.training.trainers` feeds to `Stage2OneShot` /
+`Stage2Autoregressive`. Split out of the trainers so the (target, generator,
+decoder) width rules — the fiddliest part of this codebase —
+live in one place and stay unit-testable on their own.
+"""
+
+import torch
+import torch.nn.functional as F
+
+from giant.config import ParticleTypeConfig
+from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM
+from giant.model.objectives import build_objective
+from giant.sample import sample_secondaries_ar
+
+
+def _gumbel_tau(step: int, total_steps: int, tau_start: float, tau_end: float) -> float:
+    """Linear anneal of the straight-through Gumbel-softmax temperature.
+
+    Deterministic in `step`/`total_steps` alone (no extra state), so it
+    recomputes correctly on `--resume` from a checkpoint's saved `global_step`
+    without needing to persist anything new (see
+    giant.model.network.Router.combine_weights).
+    """
+    progress = min(step / max(total_steps, 1), 1.0)
+    return tau_start + (tau_end - tau_start) * progress
+
+
+def _type_repr(
+    sec_type_idx: torch.Tensor,
+    sec_cont: torch.Tensor,
+    particle_type_cfg: ParticleTypeConfig,
+    cond_enc: torch.nn.Module,
+    emb_dim: int,
+) -> torch.Tensor:
+    """(B, K_MAX, type_dim) ground-truth type representation, generator-
+    independent (unlike `_assemble_stage2_ar_target`'s training *target*,
+    which varies by generator/objective — see its docstring): `"physical"` ->
+    `(log_mass, charge)`; `"onehot"` -> one-hot of the true class;
+    `"embedding"` -> the conditioning's own detached embedding-table row.
+
+    Used both to build `_assemble_stage2_ar_target`'s wgan+onehot/embedding
+    branch and as the AR history features' previous-secondary identity — the
+    latter must always reflect the true physical secondary that came before,
+    regardless of what the *current* token's own training objective is.
+    """
+    target = particle_type_cfg.target
+    if target == "physical":
+        return sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM]
+    if target == "onehot":
+        return F.one_hot(sec_type_idx, num_classes=emb_dim).float()
+    return cond_enc.pdg_emb(sec_type_idx).detach()
+
+
+def _assemble_stage2_ar_target(
+    sec_cont: torch.Tensor,
+    sec_type_idx: torch.Tensor,
+    particle_type_cfg: ParticleTypeConfig,
+    generator: str,
+    cond_enc: torch.nn.Module,
+    emb_dim: int,
+) -> torch.Tensor:
+    """(B, K_MAX, token_dim) ground-truth per-token target — the unflattened
+    analogue of `_assemble_stage2_real` (defined below in terms of this),
+    matching whatever width `Stage2Autoregressive`'s (or `Stage2OneShot`'s)
+    own trunk produces for this (target, generator) combination
+    (`giant.model.network.stage2_trunk_sec_dim`):
+
+    - `target = "physical"`: unchanged from v0.2 — `sec_cont` (stick_logit,
+      dir, log_mass, charge) as-is.
+    - `target` in `("onehot", "embedding")` + an objective that doesn't fold
+      the type slice (flow/ddpm): just the continuous stick/dir slots — the
+      type slice isn't part of this tensor at all (`type_head` handles it
+      separately).
+    - `target` in `("onehot", "embedding")` + a folding objective (wgan):
+      stick/dir slots concatenated with the per-slot type representation (a
+      one-hot of the true class, relaxed on the *generated* side only, by the
+      caller; or the conditioning's own detached embedding-table row).
+    """
+    target = particle_type_cfg.target
+    if target == "physical":
+        return sec_cont
+    cont = sec_cont[..., :CONT_SLOT_DIM]
+    if not build_objective(generator).folds_type_slice:
+        return cont
+    type_repr = _type_repr(sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim)
+    return torch.cat([cont, type_repr], dim=-1)
+
+
+def _assemble_stage2_real(
+    sec_cont: torch.Tensor,
+    sec_type_idx: torch.Tensor,
+    particle_type_cfg: ParticleTypeConfig,
+    generator: str,
+    cond_enc: torch.nn.Module,
+    emb_dim: int,
+) -> torch.Tensor:
+    """Ground-truth flattened stage-2 vector for `Stage2OneShot` — the
+    flattened form of `_assemble_stage2_ar_target`, which
+    `Stage2Autoregressive`'s per-token target also uses; the two must stay in
+    lockstep. See `_assemble_stage2_ar_target`'s docstring for the
+    (target, generator) width rules."""
+    return _assemble_stage2_ar_target(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim).flatten(
+        1
+    )
+
+
+def _stick_fraction(sec_cont: torch.Tensor) -> torch.Tensor:
+    """(B, K_MAX) — sigmoid of each slot's own stick-breaking logit
+    (`sec_cont[...,0]`); scale-free (see `giant.data.transforms.
+    encode_secondaries`), so this needs no absolute `e_sec`."""
+    return torch.sigmoid(sec_cont[..., 0])
+
+
+def _remaining_energy_fraction(fraction: torch.Tensor) -> torch.Tensor:
+    """(B, K_MAX) — fraction of the original e_sec budget unclaimed entering
+    slot i: `1.0` at `i=0`, `prod_{j<i}(1-fraction_j)` for `i>=1`
+    ("no re-derivation needed": the existing
+    stick-breaking encoding is already scale-free, so this is derivable from
+    the batch's ground-truth stick logits alone, no `e_sec` required).
+
+    Forced fp32 regardless of the caller's ambient `train.precision` autocast
+    region: a `cumprod` over `K_MAX` slots in bf16 underflows to zero within a
+    handful of slots, killing `remaining_frac` as a conditioning signal — the
+    numpy encoder (`giant.data.transforms.encode_secondaries`'s stick-breaking
+    twin) already promotes to float64 for exactly this reason (gitea #47)."""
+    with torch.autocast(fraction.device.type, enabled=False):
+        fraction = fraction.float()
+        cumprod = torch.cumprod(1.0 - fraction, dim=1)
+        return torch.cat([torch.ones_like(cumprod[:, :1]), cumprod[:, :-1]], dim=1)
+
+
+def _shift_prev(x: torch.Tensor) -> torch.Tensor:
+    """`(B, K, ...)` -> same shape, slot i holds slot i-1's value; slot 0 gets
+    an arbitrary zero placeholder (never read as-is — see `_ar_has_prev`;
+    `MarkovHistory` substitutes its own learned start vector there instead)."""
+    return torch.cat([torch.zeros_like(x[:, :1]), x[:, :-1]], dim=1)
+
+
+def _ar_has_prev(k_max: int, device: torch.device) -> torch.Tensor:
+    """`(1, K_MAX)` bool: True for slot index `>= 1`. Correct without
+    `n_sec`: `sec_mask` is a prefix mask, so any *valid* token at `k>=1`
+    always has a valid predecessor at `k-1`; the only wrong cases are tokens
+    that are themselves padding, already masked out of every loss."""
+    return (torch.arange(k_max, device=device) >= 1).unsqueeze(0)
+
+
+def _stop_target_and_mask(n_sec: torch.Tensor, k_max: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+    """`(target, mask)`, both `(B, K_MAX)`, for `n_sec.mode = "stop_token"`'s
+    per-slot EOS head (`Stage2Autoregressive.predict_stop`).
+
+    `predict_stop` is evaluated on slot `k`'s own (pre-token) conditioning —
+    "should generation have already stopped by here" — so `target[k] = 1`
+    exactly at `k == n_sec` (the first invalid slot: `sample_secondaries_ar`
+    checks this before spending a model call generating that slot's token),
+    `0` elsewhere. `mask` is `k <= n_sec` — one slot *wider* than
+    `StageTrainer._sec_mask`'s `k < n_sec` token-content mask, since the stop
+    slot itself (`k == n_sec`) must be supervised even though there is no
+    real secondary there. A row with `n_sec == k_max` has no in-range stop
+    slot at all: `mask` covers the full `k_max` range (every generated token
+    is real) and `target` is all-zero — `sample_secondaries_ar` correctly
+    never breaks early for it, running into the `k_max` safety cap instead."""
+    idx = torch.arange(k_max, device=device).unsqueeze(0)
+    target = (idx == n_sec.unsqueeze(1)).float()
+    mask = idx <= n_sec.unsqueeze(1)
+    return target, mask
+
+
+def _ar_meta(k_max: int, batch: int, device: torch.device, fraction: torch.Tensor) -> dict[str, torch.Tensor]:
+    """`has_prev`/`remaining_frac`/`slot_idx` — the three per-token AR
+    conditioning tensors that don't depend on *which* history representation
+    (ground truth vs. the scheduled-sampling mix) produced `fraction`.
+    Shared by `_assemble_stage2_ar_inputs` and
+    `_assemble_stage2_ar_inputs_scheduled`, which differ only in
+    `history_feat`."""
+    slot_idx = (torch.arange(k_max, device=device).float() / max(k_max - 1, 1)).unsqueeze(0)
+    return {
+        "has_prev": _ar_has_prev(k_max, device).expand(batch, -1),
+        "remaining_frac": _remaining_energy_fraction(fraction),
+        "slot_idx": slot_idx.expand(batch, -1),
+    }
+
+
+def _assemble_stage2_ar_inputs(
+    sec_cont: torch.Tensor,
+    sec_type_idx: torch.Tensor,
+    particle_type_cfg: ParticleTypeConfig,
+    cond_enc: torch.nn.Module,
+    emb_dim: int,
+) -> dict[str, torch.Tensor]:
+    """Ground-truth per-token AR conditioning tensors — all `(B, K_MAX, ...)`
+    or `(B, K_MAX)`, built in one vectorized pass (teacher forcing means
+    every token's input is ground truth).
+    Keys match `Stage2Autoregressive.forward`'s trailing kwargs."""
+    device = sec_cont.device
+    B, K = sec_cont.shape[0], sec_cont.shape[1]
+    fraction = _stick_fraction(sec_cont)
+    type_repr = _type_repr(sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim)
+    history_feat = torch.cat(
+        [
+            _shift_prev(fraction).unsqueeze(-1),
+            _shift_prev(sec_cont[..., 1:CONT_SLOT_DIM]),
+            _shift_prev(type_repr),
+        ],
+        dim=-1,
+    )
+    return {"history_feat": history_feat, **_ar_meta(K, B, device, fraction)}
+
+
+def _linear_schedule(p_start: float, p_end: float, epoch: int, total_epochs: int) -> float:
+    """Linear interpolation from `p_start` (epoch 0) to `p_end` (the final
+    epoch) — standard scheduled sampling (Bengio et al. 2015), shared by
+    every train-time schedule keyed on epoch."""
+    frac = epoch / max(total_epochs - 1, 1)
+    frac = min(max(frac, 0.0), 1.0)
+    return p_start + (p_end - p_start) * frac
+
+
+def _stage2_tf_prob(mode: str, p_start: float, p_end: float, epoch: int, total_epochs: int) -> float:
+    """P(condition slot k+1 on the TRUE token k rather than the model's own
+    prediction), for the current epoch
+    (`stage2_model.autoregressive.teacher_forcing`).
+    `"always"`/`"never"` are the two degenerate constants; `"scheduled"`
+    linearly interpolates `p_start` to `p_end` via `_linear_schedule`."""
+    if mode == "always":
+        return 1.0
+    if mode == "never":
+        return 0.0
+    return _linear_schedule(p_start, p_end, epoch, total_epochs)
+
+
+def _ctx_truth_prob(mode: str, p_start: float, p_end: float, epoch: int, total_epochs: int) -> float:
+    """P(condition stage 2 on the TRUE stage-1 outcome rather than a fresh
+    stage-1 sample), for the current epoch (`stage2_model.stage1_context`).
+    `"truth"` is the degenerate constant 1.0; `"sampled"` linearly
+    interpolates `ctx_p_start` to `ctx_p_end` via `_linear_schedule` — the
+    stage-boundary counterpart of `_stage2_tf_prob`."""
+    if mode == "truth":
+        return 1.0
+    return _linear_schedule(p_start, p_end, epoch, total_epochs)
+
+
+def _history_repr_from_ar_sample(
+    sec_cont_pred: torch.Tensor,
+    sec_type_pred: torch.Tensor,
+    particle_type_cfg: ParticleTypeConfig,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """`(fraction, direction, type_repr)` — the same triple `_type_repr` /
+    `_stick_fraction` derive from ground truth, but from a free-running
+    `sample_secondaries_ar` self-sample instead, so the two can be mixed
+    slot-by-slot under scheduled sampling (`_assemble_stage2_ar_inputs_scheduled`).
+    `target="onehot"` collapses the raw per-slot type logits to a hard
+    one-hot of `argmax` — `sample_secondaries_ar`'s own history convention
+    (see its docstring), matching what `MarkovHistory`/`AttentionHistory`
+    were trained on; the other two targets are already the right
+    representation."""
+    fraction = torch.sigmoid(sec_cont_pred[..., 0])
+    direction = sec_cont_pred[..., 1:4]
+    if particle_type_cfg.target == "onehot":
+        type_dim = sec_type_pred.size(-1)
+        type_repr = F.one_hot(sec_type_pred.argmax(-1), num_classes=type_dim).float()
+    else:
+        type_repr = sec_type_pred
+    return fraction, direction, type_repr
+
+
+def _assemble_stage2_ar_inputs_scheduled(
+    model: torch.nn.Module,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+    stage1_ctx: torch.Tensor,
+    sec_cont: torch.Tensor,
+    sec_type_idx: torch.Tensor,
+    n_sec: torch.Tensor,
+    particle_type_cfg: ParticleTypeConfig,
+    cond_enc: torch.nn.Module,
+    emb_dim: int,
+    p_tf: float,
+    sample_steps: int,
+) -> dict[str, torch.Tensor]:
+    """Scheduled-sampling counterpart of `_assemble_stage2_ar_inputs`
+    (`teacher_forcing` = "scheduled"/"never"):
+    each slot's history is the TRUE previous token with probability `p_tf`
+    (an independent per-example, per-slot Bernoulli draw) and the model's own
+    free-running prediction otherwise — closing the train/inference gap that
+    `teacher_forcing="always"` (ground truth throughout training) never sees.
+    `p_tf >= 1.0` degenerates exactly to `_assemble_stage2_ar_inputs` (and
+    skips self-sampling entirely), so callers can call this unconditionally.
+
+    The free-running estimate is a REAL autoregressive self-sample —
+    `giant.sample.sample_secondaries_ar` under `torch.no_grad()` — not a
+    cheap one-step proxy, so building it costs the same `k_max` (`* steps`
+    for flow) sequential forwards `sample.py` pays at inference, EVERY batch
+    this is called on (paid at train time too whenever teacher_forcing !=
+    "always"). Fully detached: gradient only ever flows
+    through the "real" target path each stage trainer already uses
+    (`_assemble_stage2_ar_target`), never through this self-sample.
+    """
+    device = sec_cont.device
+    B, K = sec_cont.shape[0], sec_cont.shape[1]
+    if p_tf >= 1.0:
+        return _assemble_stage2_ar_inputs(sec_cont, sec_type_idx, particle_type_cfg, cond_enc, emb_dim)
+
+    was_training = model.training
+    sec_cont_pred, sec_type_pred, _ = sample_secondaries_ar(
+        model, cond_cont, cond_cat, stage1_ctx, n_sec, steps=sample_steps
+    )
+    if was_training:
+        model.train()
+
+    fraction_gt = _stick_fraction(sec_cont)
+    dir_gt = sec_cont[..., 1:CONT_SLOT_DIM]
+    type_repr_gt = _type_repr(sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim)
+    fraction_pred, dir_pred, type_repr_pred = _history_repr_from_ar_sample(
+        sec_cont_pred, sec_type_pred, particle_type_cfg
+    )
+
+    use_gt = torch.rand(B, K, device=device) < p_tf
+    fraction = torch.where(use_gt, fraction_gt, fraction_pred)
+    direction = torch.where(use_gt.unsqueeze(-1), dir_gt, dir_pred)
+    type_repr = torch.where(use_gt.unsqueeze(-1), type_repr_gt, type_repr_pred)
+
+    own_feat = torch.cat([fraction.unsqueeze(-1), direction, type_repr], dim=-1)
+    return {
+        "history_feat": _shift_prev(own_feat),
+        **_ar_meta(K, B, device, fraction),
+    }
+
+
+def _relax_onehot_type_slice(
+    x_flat: torch.Tensor,
+    k_max: int,
+    cont_dim: int,
+    type_dim: int,
+    tau: float,
+    grad_probe: dict[str, float] | None = None,
+) -> torch.Tensor:
+    """Straight-through Gumbel-softmax relaxation of the per-slot type slice
+    inside a flattened `(B, k_max * (cont_dim + type_dim))` WGAN generator
+    output: the forward pass is a
+    hard one-hot (matching what the critic sees from real data), the
+    backward pass flows smooth gradient. Continuous slots (stick/dir, and
+    the type slice itself under `target = "embedding"`, which never calls
+    this) pass through unchanged.
+
+    `grad_probe`, if given, gets `["cont"]`/`["type"]` populated with the L2
+    norm of the gradient reaching this split point during the next
+    `.backward()` call that touches it — a backward hook, not a second
+    backward pass. This is the differentiability validation-obligation
+    instrumentation: the trunk-gradient contribution
+    from the type slice vs. the continuous slices, for
+    `particle_type.target="onehot"` + `generator="wgan"`. Only ever populated
+    on a `did_g_step` batch — the critic step backprops through
+    `fake.detach()`, which never reaches these hooks — so it stays empty
+    (callers default to `0.0`) otherwise."""
+    B = x_flat.size(0)
+    x = x_flat.view(B, k_max, cont_dim + type_dim)
+    cont, type_logits = x[..., :cont_dim], x[..., cont_dim:]
+    if grad_probe is not None:
+        cont.register_hook(lambda g: grad_probe.__setitem__("cont", g.norm().item()))
+        type_logits.register_hook(lambda g: grad_probe.__setitem__("type", g.norm().item()))
+    type_soft = F.gumbel_softmax(type_logits, tau=tau, hard=True, dim=-1)
+    return torch.cat([cont, type_soft], dim=-1).reshape(B, -1)
@@ -8,9 +8,7 @@ from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
 _SEC_PHYS_NAMES = ["log_mass", "charge"]


-def _histogram_kl(
-    p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8
-) -> float:
+def _histogram_kl(p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8) -> float:
    """KL(P || Q) between two 1D samples, estimated via a shared histogram."""
    lo = min(p_samples.min(), q_samples.min())
    hi = max(p_samples.max(), q_samples.max())
@@ -33,9 +31,7 @@ def _bincount_frac(x: np.ndarray, minlength: int) -> np.ndarray:
    return counts / total if total > 0 else counts


-def _categorical_kl(
-    real_idx: np.ndarray, gen_idx: np.ndarray, n_classes: int, eps: float = 1e-8
-) -> float:
+def _categorical_kl(real_idx: np.ndarray, gen_idx: np.ndarray, n_classes: int, eps: float = 1e-8) -> float:
    """KL(P_real || Q_gen) between two class-index samples over `n_classes`
    categories, estimated from bincount fractions. NaN if either side has no
    valid samples (mirrors `_histogram_kl`'s empty-input handling)."""
@@ -48,9 +44,7 @@ def _categorical_kl(
    return float(np.sum(p * np.log(p / q)))


-def _embedding_nearest_class(
-    vectors: torch.Tensor, emb_weight: torch.Tensor
-) -> np.ndarray:
+def _embedding_nearest_class(vectors: torch.Tensor, emb_weight: torch.Tensor) -> np.ndarray:
    """Nearest row index (L1) of `vectors` (..., emb_dim) against `emb_weight`
    (vocab, emb_dim) — same computation as
    `giant.particles.decode_embedding_nearest`, but returning the raw class
@@ -85,10 +79,10 @@ def validate_marginals(

    When `sec_decoder` is given, also validates Stage 2 via
    `giant.sample.sample_stage2`/`resolve_n_sec` (generator- and
-    one-shot-vs-autoregressive-agnostic, docs/v0.3.0-design.md §10): n_sec
+    one-shot-vs-autoregressive-agnostic): n_sec
    distribution (+ classification accuracy), per-slot energy-fraction
    marginals, and a particle-type marginal whose shape depends on
-    `sec_decoder.particle_type_cfg["target"]` — restricted to each side's own
+    `sec_decoder.particle_type_cfg.target` — restricted to each side's own
    valid slots (real: `n_sec`; generated: the resolved `n_sec_pred`), since
    the two need not agree on how many slots are valid. Adds {"n_sec_real",
    "n_sec_pred", "n_sec_accuracy", "energy_fraction_kl"} plus, under
@@ -109,11 +103,7 @@ def validate_marginals(
        sec_decoder.eval()

    k_max = sec_decoder.k_max if sec_decoder is not None else 0
-    target = (
-        sec_decoder.particle_type_cfg.get("target", "physical")
-        if sec_decoder is not None
-        else "physical"
-    )
+    target = sec_decoder.particle_type_cfg.target if sec_decoder is not None else "physical"

    all_real, all_gen = [], []
    all_n_sec_real, all_n_sec_pred = [], []
@@ -125,15 +115,12 @@ def validate_marginals(
    for i, batch in enumerate(val_loader):
        if n_batches is not None and i >= n_batches:
            break
-        # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx,
-        # sec_type_idx).
-        cond_cont, cond_cat, x1, n_sec, sec_cont, _proc_idx, sec_type_idx = batch
-        cond_cont = cond_cont.to(device)
-        cond_cat = cond_cat.to(device)
+        # batch is a StepBatch (giant.data.dataset).
+        x1, n_sec, sec_cont, sec_type_idx = batch.target_s1, batch.n_sec, batch.sec_cont, batch.sec_type_idx
+        cond_cont = batch.cond_cont.to(device)
+        cond_cat = batch.cond_cat.to(device)

-        gen, n_sec_pred = sample_stage1(
-            stage1_model, cond_cont, cond_cat, steps=steps, ddpm_steps=ddpm_steps
-        )
+        gen, n_sec_pred = sample_stage1(stage1_model, cond_cont, cond_cat, steps=steps, ddpm_steps=ddpm_steps)

        all_real.append(x1.numpy())
        all_gen.append(gen.cpu().numpy())
@@ -141,13 +128,8 @@ def validate_marginals(
        if sec_decoder is None:
            continue

-        n_sec_pred = resolve_n_sec(
-            stage1_model, sec_decoder, cond_cont, cond_cat, gen, n_sec_pred
-        )
-        n_sec_pred_np = n_sec_pred.cpu().numpy()
+        n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cond_cont, cond_cat, gen, n_sec_pred)
        n_sec_np = n_sec.numpy()
-        all_n_sec_real.append(n_sec_np)
-        all_n_sec_pred.append(n_sec_pred_np)

        real_valid = np.arange(k_max)[None, :] < n_sec_np[:, None]  # (B, k_max)
        real_frac = 1.0 / (1.0 + np.exp(-sec_cont[:, :, 0].numpy().astype(np.float64)))
@@ -155,9 +137,14 @@ def validate_marginals(
        sec_cont_pred, sec_type_pred, sec_valid_pred = sample_stage2(
            sec_decoder, cond_cont, cond_cat, gen, n_sec_pred, steps=steps
        )
-        gen_frac = 1.0 / (
-            1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64))
-        )
+        # A stop-token decoder resolves n_sec_pred=None above — read the real
+        # count back off sec_valid_pred instead (a no-op round trip under
+        # every other n_sec.mode, where sec_valid_pred was built FROM
+        # n_sec_pred in the first place).
+        n_sec_pred_np = sec_valid_pred.sum(dim=-1).cpu().numpy()
+        all_n_sec_real.append(n_sec_np)
+        all_n_sec_pred.append(n_sec_pred_np)
+        gen_frac = 1.0 / (1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64)))
        gen_valid = sec_valid_pred.cpu().numpy()

        if target == "physical":
@@ -182,17 +169,9 @@ def validate_marginals(
    real = np.concatenate(all_real, axis=0)
    generated = np.concatenate(all_gen, axis=0)

-    kl_divergence = np.array(
-        [
-            _histogram_kl(real[:, j], generated[:, j], bins=kl_bins)
-            for j in range(real.shape[1])
-        ]
-    )
+    kl_divergence = np.array([_histogram_kl(real[:, j], generated[:, j], bins=kl_bins) for j in range(real.shape[1])])

-    header = (
-        f"{'Dim':<20} {'real_mean':>10} {'gen_mean':>10} "
-        f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
-    )
+    header = f"{'Dim':<20} {'real_mean':>10} {'gen_mean':>10} {'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
    print(f"\n{header}")
    print("-" * len(header))
    for j, name in enumerate(LOCAL_TARGET_NAMES):
@@ -215,10 +194,7 @@ def validate_marginals(
    n_sec_accuracy = float((n_sec_real == n_sec_pred_all).mean())

    energy_fraction_kl = np.full(k_max, np.nan)
-    print(
-        f"\n{'n_sec':<20} accuracy={n_sec_accuracy:.4f}  "
-        f"mean|Δ|={np.abs(n_sec_real - n_sec_pred_all).mean():.4f}"
-    )
+    print(f"\n{'n_sec':<20} accuracy={n_sec_accuracy:.4f}  mean|Δ|={np.abs(n_sec_real - n_sec_pred_all).mean():.4f}")
    n_sec_dist_header = f"{'n_sec value':<20} {'real_frac':>10} {'gen_frac':>10}"
    print(n_sec_dist_header)
    print("-" * len(n_sec_dist_header))
@@ -240,10 +216,7 @@ def validate_marginals(
            continue
        kl = _histogram_kl(r, g, bins=kl_bins)
        energy_fraction_kl[j] = kl
-        print(
-            f"{j:<24} {r.mean():>10.4f} {g.mean():>10.4f} "
-            f"{r.std():>10.4f} {g.std():>10.4f} {kl:>14.4f}"
-        )
+        print(f"{j:<24} {r.mean():>10.4f} {g.mean():>10.4f} {r.std():>10.4f} {g.std():>10.4f} {kl:>14.4f}")

    result.update(
        {
@@ -259,12 +232,7 @@ def validate_marginals(
        phys_gen = np.concatenate(all_phys_gen, axis=0)  # (M, 2)

        if len(phys_real) > 0 and len(phys_gen) > 0:
-            phys_kl = np.array(
-                [
-                    _histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins)
-                    for j in range(2)
-                ]
-            )
+            phys_kl = np.array([_histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins) for j in range(2)])
        else:
            phys_kl = np.full(2, np.nan)

@@ -278,21 +246,14 @@ def validate_marginals(
            if len(r) == 0 or len(g) == 0:
                continue
            print(
-                f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} "
-                f"{r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"
+                f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} {r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"
            )

-        result.update(
-            {"phys_real": phys_real, "phys_generated": phys_gen, "phys_kl": phys_kl}
-        )
+        result.update({"phys_real": phys_real, "phys_generated": phys_gen, "phys_kl": phys_kl})
    else:
        type_class_real = np.concatenate(all_type_class_real, axis=0)
        type_class_gen = np.concatenate(all_type_class_gen, axis=0)
-        n_classes = (
-            sec_decoder.type_dim
-            if target == "onehot"
-            else sec_decoder.cond_enc.pdg_emb.weight.size(0)
-        )
+        n_classes = sec_decoder.type_dim if target == "onehot" else sec_decoder.cond_enc.pdg_emb.weight.size(0)
        type_class_kl = _categorical_kl(type_class_real, type_class_gen, n_classes)

        print(
@@ -0,0 +1,40 @@
+"""b2luigi orchestration of the full GIANT pipeline.
+
+One workflow TOML (``spec.py``) parameterises an entire experiment — dataset,
+geometry oracle, N trainings, N rollouts, N analyses — and ``giant workflow
+run <spec.toml>`` turns it into a b2luigi task graph (``tasks.py``) whose
+targets are files on ``/ceph``: nothing is recomputed that already exists,
+every step waits for its inputs, and HTCondor submission/polling is b2luigi's
+job rather than a hand-rolled submit-file generator.
+
+This is the only sanctioned way to run a multi-step pipeline; ``giant`` and
+``dwarf`` stay single-step primitives that these tasks invoke.
+
+``tasks``/``run`` import b2luigi, so they are *not* imported here — a plain
+``import giant.workflow`` (or ``giant.workflow.spec``) works without the
+``workflow`` extra installed.
+"""
+
+from giant.workflow.spec import (
+    AnalysisSpec,
+    CondorSpec,
+    DatasetSpec,
+    GeometrySpec,
+    RolloutSpec,
+    TrainSpec,
+    WorkflowSpec,
+    load_spec,
+    spec_hash,
+)
+
+__all__ = [
+    "AnalysisSpec",
+    "CondorSpec",
+    "DatasetSpec",
+    "GeometrySpec",
+    "RolloutSpec",
+    "TrainSpec",
+    "WorkflowSpec",
+    "load_spec",
+    "spec_hash",
+]
@@ -0,0 +1,82 @@
+"""HTCondor job descriptions for the workflow tasks.
+
+b2luigi writes every key of a task's ``htcondor_settings`` dict straight into
+that job's submit description, so these helpers are just the ETP-specific
+resource/requirement conventions in one place:
+
+* **CPU jobs** (setup cache, geometry oracle, analysis compute) keep what
+  the deleted ``giant analyze submit`` used: ``+RemoteJob`` for grid I/O, or
+  ``TARGET.ProvidesETPResources`` when the files are local to the cluster.
+* **GPU jobs** (training epochs, rollout) are remote-only, so they always
+  carry ``+RemoteJob`` and reach ``/ceph`` through
+  ``TARGET.ProvidesEtpCeph`` — the requirement strings are ported from the
+  ``condor-gpu-train-rollout`` branch's ``giant/condor.py`` rather than
+  rewritten, since they encode what the ETP HTCondor wiki documents for
+  TOpAS/NEMO2 GPU workers.
+"""
+
+from __future__ import annotations
+
+from giant.workflow.spec import CondorSpec
+
+__all__ = ["cpu_settings", "gpu_settings", "gpu_requirements"]
+
+
+def cpu_settings(
+    condor: CondorSpec,
+    *,
+    request_memory_mb: int | None = None,
+    request_cpus: int | None = None,
+    walltime_s: int | None = None,
+) -> dict:
+    settings: dict = {
+        "universe": "docker",
+        "docker_image": condor.docker_image_cpu,
+        "request_memory": request_memory_mb if request_memory_mb is not None else condor.request_memory_mb,
+        "request_cpus": request_cpus if request_cpus is not None else condor.request_cpus,
+        "accounting_group": condor.accounting_group,
+        "should_transfer_files": "YES",
+        "when_to_transfer_output": "ON_EXIT",
+    }
+    if condor.remote:
+        settings["+RemoteJob"] = "True"
+    else:
+        settings["requirements"] = "TARGET.ProvidesETPResources"
+    if walltime_s is not None:
+        settings["+RequestWalltime"] = int(walltime_s)
+    return settings
+
+
+def gpu_requirements(gpu_type: str | None = None, gpu_memory_mb: int | None = None) -> str:
+    """``TARGET.ProvidesEtpCeph`` (remote /ceph access) ANDed with any GPU pin."""
+    clauses = ["TARGET.ProvidesEtpCeph =?= True"]
+    if gpu_type is not None:
+        clauses.append(f'TARGET.GPUs_DeviceName =?= "{gpu_type}"')
+    if gpu_memory_mb is not None:
+        clauses.append(f"TARGET.GPUs_GlobalMemoryMb >= {gpu_memory_mb}")
+    return " && ".join(clauses)
+
+
+def gpu_settings(
+    condor: CondorSpec,
+    *,
+    request_gpus: int = 1,
+    gpu_type: str | None = None,
+    gpu_memory_mb: int | None = None,
+    request_memory_mb: int = 16384,
+    request_cpus: int = 4,
+    walltime_s: int = 86400,
+) -> dict:
+    return {
+        "universe": "docker",
+        "docker_image": condor.docker_image_gpu,
+        "request_memory": request_memory_mb,
+        "request_cpus": request_cpus,
+        "RequestGPUs": request_gpus,
+        "+RequestWalltime": int(walltime_s),
+        "accounting_group": condor.accounting_group,
+        "should_transfer_files": "YES",
+        "when_to_transfer_output": "ON_EXIT",
+        "+RemoteJob": "True",
+        "requirements": f"({gpu_requirements(gpu_type, gpu_memory_mb)})",
+    }
@@ -0,0 +1,103 @@
+#!/usr/bin/env python
+"""Entry point b2luigi re-executes on every worker.
+
+Locally this is what ``giant workflow run <spec.toml>`` execs; on a batch
+worker it is what the generated wrapper script runs (after ``cd repo_dir`` and
+sourcing ``env_script``), with ``--spec`` forwarded via the
+``task_cmd_additional_args`` setting so the worker resolves exactly the same
+spec — and therefore the same task graph and output paths — as the submitter.
+
+b2luigi needs a real script path for that re-execution, which is why this is a
+script rather than a ``python -m`` module.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+# Allow `python giant/workflow/run.py` from a checkout that isn't installed.
+sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+
+import b2luigi  # noqa: E402
+
+from giant.workflow.spec import WorkflowSpec, load_spec  # noqa: E402
+from giant.workflow.tasks import WorkflowTask, set_spec  # noqa: E402
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Run a GIANT workflow spec with b2luigi.")
+    parser.add_argument("--spec", required=True, help="Workflow TOML (see configs/workflow_example.toml)")
+    parser.add_argument("--workers", type=int, default=1, help="Concurrent luigi workers")
+    parser.add_argument(
+        "--batch",
+        action="store_true",
+        help="Submit batch-system tasks to HTCondor (otherwise everything runs locally)",
+    )
+    parser.add_argument(
+        "--mode",
+        choices=("run", "dry-run", "show-output", "remove"),
+        default="run",
+        help="run (default), dry-run (print pending tasks), show-output (print every target), remove (delete outputs)",
+    )
+    parser.add_argument("--scheduler-host", default=None, help="luigid host (default: local scheduler)")
+    parser.add_argument("--scheduler-port", type=int, default=None, help="luigid port")
+    return parser
+
+
+def configure(spec: WorkflowSpec, spec_path: Path, batch: bool) -> None:
+    """Wire b2luigi's settings from the spec.
+
+    ``/ceph`` is shared between submit host and workers, so there is
+    deliberately no ``transfer_files``: ``result_dir``/``log_dir`` must live
+    somewhere both sides can see.
+    """
+    set_spec(spec)
+
+    b2luigi.set_setting("result_dir", spec.result_dir)
+    b2luigi.set_setting("log_dir", spec.log_dir)
+    b2luigi.set_setting("task_file_dir", str(Path(spec.result_dir) / "task_files"))
+    b2luigi.set_setting("use_parameter_name_in_output", True)
+    b2luigi.set_setting("batch_system", "htcondor" if batch else "local")
+    b2luigi.set_setting("working_dir", spec.condor.repo_dir)
+    b2luigi.set_setting("job_name", spec.name)
+    if spec.condor.env_script:
+        b2luigi.set_setting("env_script", spec.condor.env_script)
+    # The worker command is `<executable> [<basename of this file>] --batch-runner
+    # --task-id ...`, run after `cd working_dir`. Only the *basename* would be
+    # used, so the filename is dropped and the repo-relative script path is
+    # made part of the executable instead.
+    b2luigi.set_setting("add_filename_to_cmd", False)
+    b2luigi.set_setting("executable", [".venv/bin/python", "giant/workflow/run.py"])
+    b2luigi.set_setting("task_cmd_additional_args", ["--spec", str(spec_path)])
+
+
+def main(argv: list[str] | None = None) -> None:
+    args, _ = build_parser().parse_known_args(argv)
+    spec_path = Path(args.spec).resolve()
+    spec = load_spec(spec_path)
+    configure(spec, spec_path, batch=args.batch)
+
+    kwargs: dict = {}
+    if args.scheduler_host:
+        kwargs["scheduler_host"] = args.scheduler_host
+    if args.scheduler_port:
+        kwargs["scheduler_port"] = args.scheduler_port
+
+    b2luigi.process(
+        WorkflowTask(workflow_name=spec.name),
+        workers=args.workers,
+        batch=args.batch,
+        dry_run=args.mode == "dry-run",
+        show_output=args.mode == "show-output",
+        remove=args.mode == "remove",
+        auto_confirm=args.mode == "remove",
+        # run.py owns --spec/--mode/...; b2luigi must not choke on them.
+        ignore_additional_command_line_args=True,
+        **kwargs,
+    )
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,383 @@
+"""Workflow TOML -> frozen dataclasses, validation, and per-task spec hashes.
+
+One spec file is the only place a pipeline is parameterised (see
+``configs/workflow_example.toml``):
+
+    [workflow]  name / result_dir / log_dir
+    [condor]    accounting group, repo dir, env script, docker images
+    [dataset]   steps (training) + reference (rollout seeds & analysis truth)
+    [geometry]  geometry-oracle build options
+    [[train]]   one per training run      (name, config, epochs, overrides, ...)
+    [[rollout]] one per rollout           (name, train = <a [[train]].name>, ...)
+    [[analysis]] one per comparison       (name, rollouts = [<[[rollout]].name>, ...])
+
+Every task carries its ``name`` plus a short ``spec_hash`` — 8 hex of the
+canonical JSON of its own resolved sub-spec **including its transitive
+parents**. That is what makes an edited spec produce a fresh result directory
+instead of silently reusing outputs computed under different settings: change
+the dataset and every hash downstream of it changes too.
+
+Unknown keys are rejected (with the valid ones listed), in the same spirit as
+``giant.config.validate_config_keys`` — a typo in a workflow spec would
+otherwise be a silently ignored setting on a multi-day pipeline.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import tomllib
+from dataclasses import MISSING, dataclass, field, fields, is_dataclass
+from pathlib import Path
+from typing import Any
+
+__all__ = [
+    "AnalysisSpec",
+    "CondorSpec",
+    "DatasetSpec",
+    "GeometrySpec",
+    "RolloutSpec",
+    "TrainSpec",
+    "WorkflowSpec",
+    "load_spec",
+    "spec_hash",
+]
+
+
+class WorkflowSpecError(ValueError):
+    """Raised for any malformed workflow spec (unknown key, bad reference, ...)."""
+
+
+# ---------------------------------------------------------------------------
+# sub-specs
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class CondorSpec:
+    """Where and how jobs run — the batch-system half of the spec.
+
+    ``repo_dir`` doubles as b2luigi's ``working_dir`` (jobs ``cd`` there before
+    running ``giant/workflow/run.py``), and ``env_script`` is sourced first,
+    since submit and worker machines don't share an environment.
+    """
+
+    accounting_group: str
+    repo_dir: str
+    env_script: str = ""
+    docker_image_cpu: str = "cverstege/alma9-gridjob"
+    docker_image_gpu: str = "mschnepf/slc7-condocker"
+    remote: bool = True
+    request_cpus: int = 1
+    request_memory_mb: int = 8192
+
+
+@dataclass(frozen=True)
+class DatasetSpec:
+    """The two datasets every pipeline needs.
+
+    ``steps`` is what training reads; ``reference`` is the held-out file
+    rollouts are seeded from and the analysis compares against (the "one
+    ground truth" premise of ``giant.analysis``).
+    """
+
+    steps: str
+    reference: str
+
+
+@dataclass(frozen=True)
+class GeometrySpec:
+    """``dwarf build-geometry-oracle`` options (see giant/tools/geometry_oracle.py)."""
+
+    method: str = "slab"
+    k: int = 1
+    subsample: int = 500_000
+    escape_factor: float = 5.0
+    seed: int = 0
+    depth_axis: int = 2
+    n_bins: int = 2000
+
+
+@dataclass(frozen=True)
+class TrainSpec:
+    """One training run, fanned out into ``ceil(epochs / epochs_per_job)`` jobs.
+
+    ``overrides`` are ``[train]``/model config keys merged on top of ``config``
+    exactly as ``giant train``'s flags are (``giant.config.merge_cli_overrides``),
+    so anything expressible on the CLI is expressible here.
+    """
+
+    name: str
+    config: str | None = None
+    epochs: int = 1
+    epochs_per_job: int = 1
+    overrides: dict[str, Any] = field(default_factory=dict)
+    request_gpus: int = 1
+    gpu_type: str | None = None
+    gpu_memory_mb: int | None = None
+    request_memory_mb: int = 16384
+    request_cpus: int = 4
+    walltime_s: int = 86400
+    num_workers: int = 4
+    shuffle_buffer: int = 65536
+    device: str | None = None
+
+
+@dataclass(frozen=True)
+class RolloutSpec:
+    """One ``giant rollout`` run against the checkpoint of ``train``."""
+
+    name: str
+    train: str
+    n_events: int | None = None
+    energy_cutoff: float = 0.1
+    max_steps: int = 1000
+    steps: int = 10
+    batch_size: int = 4096
+    max_tracks_per_event: int | None = None
+    escape_threshold: float | None = None
+    weights: str = "raw"
+    seed: int | None = None
+    request_gpus: int = 1
+    gpu_type: str | None = None
+    gpu_memory_mb: int | None = None
+    request_memory_mb: int = 16384
+    request_cpus: int = 2
+    walltime_s: int = 86400
+    device: str | None = None
+
+
+@dataclass(frozen=True)
+class AnalysisSpec:
+    """One rollout-vs-reference comparison (N rollout series, one reference)."""
+
+    name: str
+    rollouts: tuple[str, ...]
+    chunks: int = 1
+    energy_bins: int = 4
+    bins: int = 50
+    top_pdg: int = 6
+    gallery: bool = False
+    request_memory_mb: int = 8192
+    request_cpus: int = 1
+
+
+@dataclass(frozen=True)
+class WorkflowSpec:
+    """A whole pipeline: the parsed spec file plus name-keyed lookups."""
+
+    name: str
+    result_dir: str
+    log_dir: str
+    condor: CondorSpec
+    dataset: DatasetSpec
+    geometry: GeometrySpec
+    trains: tuple[TrainSpec, ...]
+    rollouts: tuple[RolloutSpec, ...]
+    analyses: tuple[AnalysisSpec, ...]
+    path: str = ""
+
+    # -- lookups ----------------------------------------------------------
+    def train(self, name: str) -> TrainSpec:
+        return _lookup(self.trains, name, "train")
+
+    def rollout(self, name: str) -> RolloutSpec:
+        return _lookup(self.rollouts, name, "rollout")
+
+    def analysis(self, name: str) -> AnalysisSpec:
+        return _lookup(self.analyses, name, "analysis")
+
+    # -- hashes -----------------------------------------------------------
+    # Each one folds in everything upstream of it, so a change anywhere in a
+    # task's ancestry moves its result directory (and only the affected
+    # subtree's).
+    def dataset_hash(self) -> str:
+        return spec_hash(self.dataset)
+
+    def warm_cache_hash(self, train_name: str) -> str:
+        # The setup cache depends on the dataset and on what this training's
+        # config asks of it (val split, conditioning, router) — not on how
+        # many epochs it runs for, so epochs/resources are deliberately left
+        # out and two trainings sharing a config share one warm-cache job.
+        t = self.train(train_name)
+        return spec_hash(self.dataset, t.config, t.overrides)
+
+    def geometry_hash(self) -> str:
+        return spec_hash(self.dataset, self.geometry)
+
+    def train_hash(self, name: str) -> str:
+        return spec_hash(self.dataset, self.train(name))
+
+    def rollout_hash(self, name: str) -> str:
+        ro = self.rollout(name)
+        return spec_hash(self.dataset, self.geometry, self.train(ro.train), ro)
+
+    def analysis_hash(self, name: str) -> str:
+        an = self.analysis(name)
+        parents = [self.rollout(r) for r in an.rollouts]
+        train_parents = [self.train(r.train) for r in parents]
+        return spec_hash(self.dataset, self.geometry, train_parents, parents, an)
+
+
+def _lookup(items, name: str, kind: str):
+    for item in items:
+        if item.name == name:
+            return item
+    known = ", ".join(sorted(i.name for i in items)) or "(none defined)"
+    raise WorkflowSpecError(f"no [[{kind}]] named {name!r} in this workflow — defined: {known}")
+
+
+# ---------------------------------------------------------------------------
+# hashing
+# ---------------------------------------------------------------------------
+
+
+def spec_hash(*parts: Any, length: int = 8) -> str:
+    """Short stable hash of one or more (sub-)specs.
+
+    Canonical JSON (sorted keys, dataclasses expanded) so the value depends
+    only on the resolved settings — not on key order in the TOML, nor on
+    which defaults were written out explicitly.
+    """
+    payload = json.dumps([_canonical(p) for p in parts], sort_keys=True, separators=(",", ":"))
+    return hashlib.sha256(payload.encode()).hexdigest()[:length]
+
+
+def _canonical(value: Any) -> Any:
+    if is_dataclass(value) and not isinstance(value, type):
+        return {f.name: _canonical(getattr(value, f.name)) for f in fields(value)}
+    if isinstance(value, dict):
+        return {str(k): _canonical(v) for k, v in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_canonical(v) for v in value]
+    if isinstance(value, Path):
+        return str(value)
+    return value
+
+
+# ---------------------------------------------------------------------------
+# parsing
+# ---------------------------------------------------------------------------
+
+
+def _build(cls, data: dict, where: str):
+    """Instantiate a frozen sub-spec, rejecting unknown/missing keys loudly."""
+    valid = {f.name for f in fields(cls)}
+    unknown = sorted(set(data) - valid)
+    if unknown:
+        raise WorkflowSpecError(f"{where}: unknown key(s) {unknown} — valid keys: {sorted(valid)}")
+    required = {f.name for f in fields(cls) if f.default is MISSING and f.default_factory is MISSING}
+    missing = sorted(required - set(data))
+    if missing:
+        raise WorkflowSpecError(f"{where}: missing required key(s) {missing}")
+    return cls(**data)
+
+
+def load_spec(path: str | Path) -> WorkflowSpec:
+    """Parse and validate a workflow TOML file."""
+    path = Path(path)
+    try:
+        raw = tomllib.loads(path.read_text())
+    except tomllib.TOMLDecodeError as exc:
+        raise WorkflowSpecError(f"{path}: not valid TOML — {exc}") from exc
+    return parse_spec(raw, path=path)
+
+
+def parse_spec(raw: dict, path: str | Path = "") -> WorkflowSpec:
+    """Validate an already-parsed workflow spec mapping."""
+    top_valid = {"workflow", "condor", "dataset", "geometry", "train", "rollout", "analysis"}
+    unknown = sorted(set(raw) - top_valid)
+    if unknown:
+        raise WorkflowSpecError(
+            f"{path or '<spec>'}: unknown top-level table(s) {unknown} — valid: {sorted(top_valid)}"
+        )
+
+    for required in ("workflow", "condor", "dataset"):
+        if required not in raw:
+            raise WorkflowSpecError(f"{path or '<spec>'}: missing required [{required}] table")
+
+    wf = dict(raw["workflow"])
+    wf_valid = {"name", "result_dir", "log_dir"}
+    wf_unknown = sorted(set(wf) - wf_valid)
+    if wf_unknown:
+        raise WorkflowSpecError(f"[workflow]: unknown key(s) {wf_unknown} — valid keys: {sorted(wf_valid)}")
+    if "name" not in wf or "result_dir" not in wf:
+        raise WorkflowSpecError("[workflow]: 'name' and 'result_dir' are required")
+    result_dir = str(Path(wf["result_dir"]).expanduser())
+    log_dir = str(Path(wf.get("log_dir", Path(result_dir) / "logs")).expanduser())
+
+    condor = _build(CondorSpec, dict(raw["condor"]), "[condor]")
+    dataset = _build(DatasetSpec, dict(raw["dataset"]), "[dataset]")
+    geometry = _build(GeometrySpec, dict(raw.get("geometry", {})), "[geometry]")
+
+    trains = tuple(_build(TrainSpec, dict(t), f"[[train]] #{i}") for i, t in enumerate(raw.get("train", [])))
+    rollouts = tuple(_build(RolloutSpec, dict(r), f"[[rollout]] #{i}") for i, r in enumerate(raw.get("rollout", [])))
+    analyses = tuple(
+        _build(AnalysisSpec, {**a, "rollouts": tuple(a.get("rollouts", ()))}, f"[[analysis]] #{i}")
+        for i, a in enumerate(raw.get("analysis", []))
+    )
+
+    _check_unique(trains, "train")
+    _check_unique(rollouts, "rollout")
+    _check_unique(analyses, "analysis")
+
+    train_names = {t.name for t in trains}
+    for ro in rollouts:
+        if ro.train not in train_names:
+            raise WorkflowSpecError(
+                f"[[rollout]] {ro.name!r}: train={ro.train!r} names no [[train]] — defined: {sorted(train_names)}"
+            )
+    rollout_names = {r.name for r in rollouts}
+    for an in analyses:
+        if not an.rollouts:
+            raise WorkflowSpecError(f"[[analysis]] {an.name!r}: 'rollouts' must name at least one [[rollout]]")
+        for r in an.rollouts:
+            if r not in rollout_names:
+                raise WorkflowSpecError(
+                    f"[[analysis]] {an.name!r}: rollout {r!r} is not defined — "
+                    f"defined: {sorted(rollout_names) or '(none)'}"
+                )
+        if len(set(an.rollouts)) != len(an.rollouts):
+            raise WorkflowSpecError(f"[[analysis]] {an.name!r}: repeated rollout name(s) in 'rollouts'")
+        if an.chunks < 1:
+            raise WorkflowSpecError(f"[[analysis]] {an.name!r}: chunks must be >= 1, got {an.chunks}")
+
+    for t in trains:
+        if t.epochs < 1:
+            raise WorkflowSpecError(f"[[train]] {t.name!r}: epochs must be >= 1, got {t.epochs}")
+        if t.epochs_per_job < 1:
+            raise WorkflowSpecError(f"[[train]] {t.name!r}: epochs_per_job must be >= 1, got {t.epochs_per_job}")
+
+    return WorkflowSpec(
+        name=wf["name"],
+        result_dir=result_dir,
+        log_dir=log_dir,
+        condor=condor,
+        dataset=dataset,
+        geometry=geometry,
+        trains=trains,
+        rollouts=rollouts,
+        analyses=analyses,
+        path=str(path),
+    )
+
+
+def _check_unique(items, kind: str) -> None:
+    names = [i.name for i in items]
+    dupes = sorted({n for n in names if names.count(n) > 1})
+    if dupes:
+        raise WorkflowSpecError(f"[[{kind}]] names must be unique — repeated: {dupes}")
+
+
+def epoch_milestones(train: TrainSpec) -> list[int]:
+    """Cumulative epoch counts, one per chained ``TrainEpochTask``.
+
+    ``epochs_per_job`` trades queue waits against job length: with
+    ``epochs=10, epochs_per_job=3`` this is ``[3, 6, 9, 10]``, i.e. job *k*
+    resumes job *k-1*'s ``last.pt`` and trains up to its own milestone.
+    """
+    step = train.epochs_per_job
+    milestones = list(range(step, train.epochs + 1, step))
+    if not milestones or milestones[-1] != train.epochs:
+        milestones.append(train.epochs)
+    return milestones
@@ -0,0 +1,647 @@
+"""The b2luigi task graph: cache-warm -> train -> rollout -> analysis.
+
+    DatasetTask (external)  ─┬─> WarmCacheTask(train) ─> TrainEpochTask(train, 1..N) ─> TrainTask(train) ─┐
+                             └─> GeometryOracleTask ──┐                                                   │
+                                                      └──> RolloutTask(rollout) <────────────────────────┘
+                                                                       │
+    AnalysisPrepTask(analysis) ─> AnalysisComputeTask(analysis, plot, chunk) ─> AnalysisRenderTask(analysis)
+                                                                                        ^
+    WorkflowTask (wrapper) ─────────────────────────────────────────────────────────────┘
+
+Every task's output directory is ``<result_dir>/<kind>/name=<name>/spec_hash=
+<hash>/…`` — the hash covers the task's resolved sub-spec *and its transitive
+parents* (``giant/workflow/spec.py``), so editing the spec produces a fresh
+directory for exactly the affected subtree instead of silently reusing stale
+outputs.
+
+Task bodies never reimplement anything: they call the same entry points the
+CLIs do (``run_warm_setup_cache``, ``run_build_geometry_oracle``,
+``run_train_job``, ``giant.analysis.prep``/``compute_one``/``merge_all``,
+``render_run``), or shell out to ``giant rollout``, which has no library-level
+entry point of its own.
+
+Training is fanned out into **one short GPU job per epoch** (or per
+``epochs_per_job`` epochs): job *k* runs ``run_train_job`` with ``epochs = k``
+and ``resume = <job k-1>/last.pt``, which the training loop already handles
+(``giant/training/loop.py`` sets ``start_epoch = ckpt["epoch"] + 1`` and
+returns early when the checkpoint already covers ``epochs``). A 200-epoch run
+then becomes 200 schedulable jobs that survive preemption and give luigi a
+real progress signal, at the cost of one (cache-warmed) setup scan and one
+queue wait per job.
+"""
+
+from __future__ import annotations
+
+import json
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+import b2luigi
+
+from giant.workflow.htcondor import cpu_settings, gpu_settings
+from giant.workflow.spec import WorkflowSpec, epoch_milestones
+
+__all__ = [
+    "AnalysisComputeTask",
+    "AnalysisPrepTask",
+    "AnalysisRenderTask",
+    "DatasetTask",
+    "GeometryOracleTask",
+    "RolloutTask",
+    "TrainEpochTask",
+    "TrainTask",
+    "WarmCacheTask",
+    "WorkflowTask",
+    "analysis_dir",
+    "analysis_jobs",
+    "get_spec",
+    "set_spec",
+]
+
+
+# ---------------------------------------------------------------------------
+# the active spec
+# ---------------------------------------------------------------------------
+# luigi parameters must be simple scalars, so tasks carry only `name` +
+# `spec_hash` and read the rest out of the one spec this process was started
+# with. Batch workers re-execute `run.py --spec <same file>` (see
+# `task_cmd_additional_args` there), so they resolve the identical spec.
+
+_SPEC: WorkflowSpec | None = None
+
+
+def set_spec(spec: WorkflowSpec) -> None:
+    global _SPEC
+    _SPEC = spec
+
+
+def get_spec() -> WorkflowSpec:
+    if _SPEC is None:
+        raise RuntimeError("no workflow spec loaded — call giant.workflow.tasks.set_spec() first")
+    return _SPEC
+
+
+def _result_dir(*parts: str) -> Path:
+    return Path(get_spec().result_dir).joinpath(*parts)
+
+
+def _task_dir(kind: str, name: str, spec_hash: str) -> Path:
+    """``<result_dir>/<kind>/name=<name>/spec_hash=<hash>``."""
+    return _result_dir(kind, f"name={name}", f"spec_hash={spec_hash}")
+
+
+def analysis_dir(spec: WorkflowSpec, name: str) -> Path:
+    """The analysis run directory — what ``prep`` lays out and every later step reads."""
+    return Path(spec.result_dir) / "analysis" / f"name={name}" / f"spec_hash={spec.analysis_hash(name)}"
+
+
+def analysis_jobs(spec: WorkflowSpec, name: str) -> list[tuple[str, int]]:
+    """Every ``(plot_id, chunk)`` compute job of one analysis.
+
+    ``chunkable=False`` specs (the checkpoint-bound diagnostics, already
+    bounded/subsampled) always run as a single chunk — the same rule the
+    deleted ``_job_walltimes`` applied.
+    """
+    from giant.analysis.catalog import catalog_ids, get_spec as get_plot_spec
+
+    chunks = spec.analysis(name).chunks
+    jobs: list[tuple[str, int]] = []
+    for plot_id in catalog_ids():
+        n = chunks if get_plot_spec(plot_id).chunkable else 1
+        jobs.extend((plot_id, chunk) for chunk in range(n))
+    return jobs
+
+
+def _giant_cmd() -> list[str]:
+    """How to invoke the ``giant`` CLI from inside a task (worker or locally)."""
+    return [sys.executable, "-m", "giant.cli"]
+
+
+# ---------------------------------------------------------------------------
+# inputs
+# ---------------------------------------------------------------------------
+
+
+class DatasetTask(b2luigi.ExternalTask):
+    """A steps parquet file or directory that must already exist.
+
+    Nothing produces it, so a missing path is a hard, immediate error rather
+    than a job that fails hours later — the usual cause being ``/ceph`` not
+    mounted on the machine the workflow was started from.
+    """
+
+    path = b2luigi.Parameter()
+
+    def output(self):
+        return b2luigi.LocalTarget(str(self.path))
+
+    def complete(self):
+        if not Path(str(self.path)).exists():
+            raise FileNotFoundError(
+                f"dataset {self.path!r} does not exist — is /ceph mounted on this machine? "
+                "(see CLAUDE.md's Compute environment section)"
+            )
+        return True
+
+
+# ---------------------------------------------------------------------------
+# setup stage
+# ---------------------------------------------------------------------------
+
+
+class WarmCacheTask(b2luigi.Task):
+    """Precompute one training's setup-stage sidecar (vocab maps, event split,
+    normalizer stats) so every per-epoch job is a cache hit instead of a
+    full rescan.
+
+    The real product (``<data>.giant_train_cache.json``) lives next to the
+    dataset, not under ``result_dir``, so the target here is a small stamp
+    recording that sidecar's path/mtime/size.
+    """
+
+    name = b2luigi.Parameter()
+    spec_hash = b2luigi.Parameter()
+
+    @property
+    def htcondor_settings(self):
+        spec = get_spec()
+        return cpu_settings(spec.condor, request_memory_mb=32768, request_cpus=4, walltime_s=21600)
+
+    def requires(self):
+        yield DatasetTask(path=get_spec().dataset.steps)
+
+    def output(self):
+        return b2luigi.LocalTarget(str(_task_dir("warm_cache", str(self.name), str(self.spec_hash)) / "stamp.json"))
+
+    def run(self):
+        from giant.data.setup_cache import sidecar_path
+        from giant.tools.warm_setup_cache import run_warm_setup_cache
+
+        spec = get_spec()
+        train = spec.train(str(self.name))
+        run_warm_setup_cache(
+            data=spec.dataset.steps,
+            config_path=Path(train.config) if train.config else None,
+        )
+        sidecar = Path(sidecar_path(spec.dataset.steps))
+        stamp = {
+            "sidecar": str(sidecar),
+            "mtime": sidecar.stat().st_mtime if sidecar.exists() else None,
+            "size": sidecar.stat().st_size if sidecar.exists() else None,
+        }
+        out = Path(self.output().path)
+        out.parent.mkdir(parents=True, exist_ok=True)
+        out.write_text(json.dumps(stamp, indent=2))
+
+
+class GeometryOracleTask(b2luigi.Task):
+    """Build the position -> (material, layer_id) oracle every rollout needs."""
+
+    spec_hash = b2luigi.Parameter()
+
+    @property
+    def htcondor_settings(self):
+        spec = get_spec()
+        return cpu_settings(spec.condor, request_memory_mb=32768, request_cpus=4, walltime_s=21600)
+
+    def requires(self):
+        yield DatasetTask(path=get_spec().dataset.steps)
+
+    def output(self):
+        return b2luigi.LocalTarget(
+            str(_result_dir("geometry", f"spec_hash={self.spec_hash}") / "oracle.pkl"),
+        )
+
+    def run(self):
+        from giant.tools.geometry_oracle import run_build_geometry_oracle
+
+        spec = get_spec()
+        g = spec.geometry
+        out = Path(self.output().path)
+        out.parent.mkdir(parents=True, exist_ok=True)
+        run_build_geometry_oracle(
+            data=Path(spec.dataset.steps),
+            out=out,
+            method=g.method,
+            k=g.k,
+            subsample=g.subsample,
+            escape_factor=g.escape_factor,
+            seed=g.seed,
+            depth_axis=g.depth_axis,
+            n_bins=g.n_bins,
+        )
+
+
+# ---------------------------------------------------------------------------
+# training
+# ---------------------------------------------------------------------------
+
+
+def _train_cfg(spec: WorkflowSpec, name: str, epochs: int) -> dict:
+    """The merged config one training job runs, resolved exactly as `giant train` does."""
+    from giant import config as gconfig
+
+    train = spec.train(name)
+    flags = {**train.overrides, "epochs": epochs}
+    overrides = gconfig.overrides_from_flags(flags)
+    cfg = gconfig.merge_cli_overrides(
+        gconfig.DEFAULT_CONFIG,
+        Path(train.config) if train.config else None,
+        overrides,
+    )
+    gconfig.validate_config(cfg, resume=True)
+    return cfg
+
+
+class TrainEpochTask(b2luigi.Task):
+    """Epochs up to ``milestone`` of one training, resuming the previous job.
+
+    Target is ``last.pt``. ``best.pt`` is written by the loop *only when that
+    epoch improved*, and ``best_val_loss`` travels inside the checkpoint, so
+    the global best comparison stays correct across jobs: "``best.pt`` exists
+    in milestone dir *k*" means exactly "one of that job's epochs was the best
+    so far".
+    """
+
+    name = b2luigi.Parameter()
+    spec_hash = b2luigi.Parameter()
+    milestone = b2luigi.IntParameter()
+
+    @property
+    def htcondor_settings(self):
+        spec = get_spec()
+        train = spec.train(str(self.name))
+        return gpu_settings(
+            spec.condor,
+            request_gpus=train.request_gpus,
+            gpu_type=train.gpu_type,
+            gpu_memory_mb=train.gpu_memory_mb,
+            request_memory_mb=train.request_memory_mb,
+            request_cpus=train.request_cpus,
+            walltime_s=train.walltime_s,
+        )
+
+    @property
+    def _dir(self) -> Path:
+        return _task_dir("train_epoch", str(self.name), str(self.spec_hash)) / f"epochs={int(self.milestone)}"
+
+    def _previous_milestone(self) -> int | None:
+        spec = get_spec()
+        milestones = epoch_milestones(spec.train(str(self.name)))
+        index = milestones.index(int(self.milestone))
+        return milestones[index - 1] if index > 0 else None
+
+    def requires(self):
+        previous = self._previous_milestone()
+        if previous is None:
+            yield WarmCacheTask(name=self.name, spec_hash=get_spec().warm_cache_hash(str(self.name)))
+        else:
+            yield TrainEpochTask(name=self.name, spec_hash=self.spec_hash, milestone=previous)
+
+    def output(self):
+        return b2luigi.LocalTarget(str(self._dir / "last.pt"))
+
+    def run(self):
+        import torch
+
+        from giant import config as gconfig
+        from giant.pipeline import run_train_job
+
+        spec = get_spec()
+        train = spec.train(str(self.name))
+        cfg = _train_cfg(spec, str(self.name), int(self.milestone))
+
+        previous = self._previous_milestone()
+        resume = None
+        if previous is not None:
+            resume = _task_dir("train_epoch", str(self.name), str(self.spec_hash)) / f"epochs={previous}" / "last.pt"
+
+        device = torch.device(train.device) if train.device else gconfig.auto_device()
+        out_dir = self._dir
+        out_dir.mkdir(parents=True, exist_ok=True)
+        run_train_job(
+            data=Path(spec.dataset.steps),
+            cfg=cfg,
+            out_dir=out_dir,
+            device=device,
+            shuffle_buffer=train.shuffle_buffer,
+            num_workers=train.num_workers,
+            resume=resume,
+            cache_setup=True,
+        )
+
+
+class TrainTask(b2luigi.Task):
+    """Publish one training's canonical outputs, hiding the epoch fan-out.
+
+    Everything downstream (``RolloutTask``, humans, ``giant analyze metrics``)
+    points here and never has to know which milestone directory happened to
+    hold the best checkpoint.
+    """
+
+    name = b2luigi.Parameter()
+    spec_hash = b2luigi.Parameter()
+    batch_system = "local"
+
+    @property
+    def _milestones(self) -> list[int]:
+        return epoch_milestones(get_spec().train(str(self.name)))
+
+    def requires(self):
+        yield TrainEpochTask(name=self.name, spec_hash=self.spec_hash, milestone=self._milestones[-1])
+
+    @property
+    def _dir(self) -> Path:
+        return _task_dir("train", str(self.name), str(self.spec_hash))
+
+    def output(self):
+        d = self._dir
+        return {
+            "best.pt": b2luigi.LocalTarget(str(d / "best.pt")),
+            "last.pt": b2luigi.LocalTarget(str(d / "last.pt")),
+            "metrics.csv": b2luigi.LocalTarget(str(d / "metrics.csv")),
+        }
+
+    def run(self):
+        epoch_base = _task_dir("train_epoch", str(self.name), str(self.spec_hash))
+        milestone_dirs = [epoch_base / f"epochs={m}" for m in self._milestones]
+
+        best_dirs = [d for d in milestone_dirs if (d / "best.pt").exists()]
+        if not best_dirs:
+            raise FileNotFoundError(
+                f"no best.pt in any milestone directory under {epoch_base} — "
+                "did every epoch job run with a validation split?"
+            )
+        out = self._dir
+        out.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(best_dirs[-1] / "best.pt", out / "best.pt")
+        shutil.copy2(milestone_dirs[-1] / "last.pt", out / "last.pt")
+        for extra in ("config.toml", "run_meta.json"):
+            src = milestone_dirs[-1] / extra
+            if src.exists():
+                shutil.copy2(src, out / extra)
+
+        # One metrics.csv for the whole run: the first job's header, then
+        # every job's rows in epoch order, so `giant analyze metrics` sees a
+        # single continuous training curve.
+        lines: list[str] = []
+        header: str | None = None
+        for d in milestone_dirs:
+            csv = d / "metrics.csv"
+            if not csv.exists():
+                continue
+            rows = csv.read_text().splitlines()
+            if not rows:
+                continue
+            if header is None:
+                header = rows[0]
+            lines.extend(rows[1:])
+        (out / "metrics.csv").write_text("\n".join([header or ""] + lines) + "\n")
+
+
+# ---------------------------------------------------------------------------
+# rollout
+# ---------------------------------------------------------------------------
+
+
+class RolloutTask(b2luigi.Task):
+    """Roll one trained checkpoint forward into full showers.
+
+    ``giant rollout`` has no library-level entry point, so this shells out to
+    the CLI — with an explicit ``--out``, which puts the YAML sidecar at the
+    deterministic ``rollout.yaml`` next to the parquet (see
+    ``giant/cli.py:_write_prediction_ref``).
+    """
+
+    name = b2luigi.Parameter()
+    spec_hash = b2luigi.Parameter()
+
+    @property
+    def htcondor_settings(self):
+        spec = get_spec()
+        ro = spec.rollout(str(self.name))
+        return gpu_settings(
+            spec.condor,
+            request_gpus=ro.request_gpus,
+            gpu_type=ro.gpu_type,
+            gpu_memory_mb=ro.gpu_memory_mb,
+            request_memory_mb=ro.request_memory_mb,
+            request_cpus=ro.request_cpus,
+            walltime_s=ro.walltime_s,
+        )
+
+    @property
+    def _dir(self) -> Path:
+        return _task_dir("rollout", str(self.name), str(self.spec_hash))
+
+    def requires(self):
+        spec = get_spec()
+        ro = spec.rollout(str(self.name))
+        yield TrainTask(name=ro.train, spec_hash=spec.train_hash(ro.train))
+        yield GeometryOracleTask(spec_hash=spec.geometry_hash())
+        yield DatasetTask(path=spec.dataset.reference)
+
+    def output(self):
+        d = self._dir
+        return {
+            "rollout.parquet": b2luigi.LocalTarget(str(d / "rollout.parquet")),
+            "rollout.yaml": b2luigi.LocalTarget(str(d / "rollout.yaml")),
+        }
+
+    def run(self):
+        spec = get_spec()
+        ro = spec.rollout(str(self.name))
+        out = self._dir / "rollout.parquet"
+        out.parent.mkdir(parents=True, exist_ok=True)
+
+        checkpoint = _task_dir("train", ro.train, spec.train_hash(ro.train)) / "best.pt"
+        oracle = _result_dir("geometry", f"spec_hash={spec.geometry_hash()}") / "oracle.pkl"
+
+        cmd = [
+            *_giant_cmd(),
+            "rollout",
+            spec.dataset.reference,
+            "--checkpoint",
+            str(checkpoint),
+            "--geometry",
+            str(oracle),
+            "--out",
+            str(out),
+            "--energy-cutoff",
+            str(ro.energy_cutoff),
+            "--max-steps",
+            str(ro.max_steps),
+            "--steps",
+            str(ro.steps),
+            "--batch-size",
+            str(ro.batch_size),
+            "--weights",
+            ro.weights,
+        ]
+        for flag, value in (
+            ("--n-events", ro.n_events),
+            ("--max-tracks-per-event", ro.max_tracks_per_event),
+            ("--escape-threshold", ro.escape_threshold),
+            ("--seed", ro.seed),
+            ("--device", ro.device),
+        ):
+            if value is not None:
+                cmd += [flag, str(value)]
+        subprocess.run(cmd, check=True)
+
+
+# ---------------------------------------------------------------------------
+# analysis
+# ---------------------------------------------------------------------------
+
+
+class AnalysisPrepTask(b2luigi.Task):
+    """Resolve the shared bin edges/group sets once, for every compute job.
+
+    Cheap and streaming, so it runs locally: everything after it needs
+    ``shared.json``/``run_meta.json`` to already exist.
+    """
+
+    name = b2luigi.Parameter()
+    spec_hash = b2luigi.Parameter()
+    batch_system = "local"
+
+    def requires(self):
+        spec = get_spec()
+        for rollout_name in spec.analysis(str(self.name)).rollouts:
+            yield RolloutTask(name=rollout_name, spec_hash=spec.rollout_hash(rollout_name))
+
+    @property
+    def _dir(self) -> Path:
+        return analysis_dir(get_spec(), str(self.name))
+
+    def output(self):
+        d = self._dir
+        return {
+            "shared.json": b2luigi.LocalTarget(str(d / "shared.json")),
+            "run_meta.json": b2luigi.LocalTarget(str(d / "run_meta.json")),
+        }
+
+    def run(self):
+        from giant.analysis import prep
+
+        spec = get_spec()
+        an = spec.analysis(str(self.name))
+        yamls = [_task_dir("rollout", r, spec.rollout_hash(r)) / "rollout.yaml" for r in an.rollouts]
+        prep(
+            yamls,
+            run_dir=self._dir,
+            n_chunks=an.chunks,
+            labels=list(an.rollouts),
+            n_energy_bins=an.energy_bins,
+            n_marginal_bins=an.bins,
+            top_k_pdg=an.top_pdg,
+        )
+
+
+class AnalysisComputeTask(b2luigi.Task):
+    """One (plot, chunk) streaming reduction — the replaced ``jobs.txt`` row.
+
+    The output path is the on-disk contract ``compute-one``/``merge_one``
+    already share (``reduced_partial/<id>__<chunk>.json``), declared
+    explicitly rather than through b2luigi's own output naming so that
+    contract is untouched.
+    """
+
+    name = b2luigi.Parameter()
+    spec_hash = b2luigi.Parameter()
+    plot_id = b2luigi.Parameter()
+    chunk = b2luigi.IntParameter()
+
+    @property
+    def htcondor_settings(self):
+        # A property, so it is evaluated at submit time — i.e. after prep has
+        # written run_meta.json, whose row counts size the walltime request.
+        from giant.analysis import RunMeta
+        from giant.analysis.runtime_estimate import estimate_runtime_s
+
+        spec = get_spec()
+        an = spec.analysis(str(self.name))
+        walltime = None
+        meta_path = analysis_dir(spec, str(self.name)) / "run_meta.json"
+        if meta_path.exists():
+            from giant.analysis.catalog import get_spec as get_plot_spec
+
+            meta = RunMeta.load(meta_path)
+            chunkable = get_plot_spec(str(self.plot_id)).chunkable
+            n_rows = meta.rows_per_chunk[int(self.chunk)] if chunkable and meta.rows_per_chunk else meta.total_rows
+            walltime = estimate_runtime_s(str(self.plot_id), n_rows)
+        return cpu_settings(
+            spec.condor,
+            request_memory_mb=an.request_memory_mb,
+            request_cpus=an.request_cpus,
+            walltime_s=walltime,
+        )
+
+    def requires(self):
+        yield AnalysisPrepTask(name=self.name, spec_hash=self.spec_hash)
+
+    def output(self):
+        run_dir = analysis_dir(get_spec(), str(self.name))
+        return b2luigi.LocalTarget(str(run_dir / "reduced_partial" / f"{self.plot_id}__{int(self.chunk)}.json"))
+
+    def run(self):
+        from giant.analysis import compute_one
+
+        compute_one(str(self.plot_id), analysis_dir(get_spec(), str(self.name)), chunk_index=int(self.chunk))
+
+
+class AnalysisRenderTask(b2luigi.Task):
+    """Merge every plot's chunk partials, then render the PDFs + gallery.
+
+    Always local — this is the only step that imports plotstyle/LaTeX, which
+    the compute worker images don't have.
+    """
+
+    name = b2luigi.Parameter()
+    spec_hash = b2luigi.Parameter()
+    batch_system = "local"
+
+    def requires(self):
+        spec = get_spec()
+        for plot_id, chunk in analysis_jobs(spec, str(self.name)):
+            yield AnalysisComputeTask(
+                name=self.name,
+                spec_hash=self.spec_hash,
+                plot_id=plot_id,
+                chunk=chunk,
+            )
+
+    def output(self):
+        run_dir = analysis_dir(get_spec(), str(self.name))
+        return b2luigi.LocalTarget(str(run_dir / "plots" / "metadata.yaml"))
+
+    def run(self):
+        # render_run joins every plot's chunk partials (merge_all) before
+        # rendering, so this one call is the whole merge+render step.
+        from giant.analysis.render import render_run
+
+        spec = get_spec()
+        render_run(analysis_dir(spec, str(self.name)), run_gallery=spec.analysis(str(self.name)).gallery)
+
+
+class WorkflowTask(b2luigi.WrapperTask):
+    """The whole pipeline: every analysis in the spec, rendered."""
+
+    workflow_name = b2luigi.Parameter()
+
+    def requires(self):
+        spec = get_spec()
+        if not spec.analyses:
+            # A spec with no [[analysis]] still has work to do — fall back to
+            # the deepest tasks it does define.
+            for ro in spec.rollouts:
+                yield RolloutTask(name=ro.name, spec_hash=spec.rollout_hash(ro.name))
+            if not spec.rollouts:
+                for tr in spec.trains:
+                    yield TrainTask(name=tr.name, spec_hash=spec.train_hash(tr.name))
+            return
+        for an in spec.analyses:
+            yield AnalysisRenderTask(name=an.name, spec_hash=spec.analysis_hash(an.name))
@@ -1,6 +1,6 @@
 [project]
 name = "giant"
-version = "0.2.0"
+version = "0.3.10"
 description = "Geant4 step-function surrogate via conditional flow matching"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -23,9 +23,12 @@ cuda = [
 ]
 dev = [
    "pytest>=8,<10",
+    "pytest-cov>=5,<8",
    "ruff>=0.15,<1",
    "ty>=0.0.50,<0.1",
-    "giant[convert,analysis,geometry,wandb]",
+    "bump-my-version>=1.2,<2",
+    "git-cliff>=2,<3",
+    "giant[convert,analysis,geometry,wandb,workflow]",
 ]
 geometry = [
    "scikit-learn>=1.4,<2",
@@ -46,17 +49,35 @@ analysis = [
    # `giant analyze render` step imports it; compute workers never do.
    "plotstyle>=1.0.0",
 ]
+# b2luigi pulls luigi + tenacity; the only sanctioned way to chain a
+# multi-step pipeline (see giant/workflow/).
+workflow = [
+    "b2luigi>=1.0,<2",
+]

 [project.scripts]
 giant = "giant.cli:app"
-dwarf = "scripts.dwarf:app"
+dwarf = "giant.tools.dwarf:app"
+
+[tool.ruff]
+line-length = 120
+
+[tool.coverage.run]
+source = ["giant"]
+omit = ["*/legacy/*"]
+
+[tool.coverage.report]
+exclude_also = [
+    "if TYPE_CHECKING:",
+    "raise NotImplementedError",
+]

 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"

 [tool.hatch.build.targets.wheel]
-packages = ["giant", "scripts"]
+packages = ["giant"]

 [tool.uv]
 conflicts = [
@@ -87,3 +108,12 @@ explicit = true
 name = "larsbogner"
 url = "https://git.larsbogner.de/api/packages/lars/pypi/simple/"
 explicit = true
+
+# luigi builds task constructors from class-level Parameter descriptors, so a
+# static checker sees no keyword parameters at all on `Task(name=..., ...)`.
+# The workflow code is written against that API; nothing else in the repo is.
+[[tool.ty.overrides]]
+include = ["giant/workflow/**", "tests/test_workflow_tasks.py"]
+
+[tool.ty.overrides.rules]
+unknown-argument = "ignore"
@@ -1,68 +0,0 @@
-"""dwarf warm-cache — precompute `giant train`'s setup-stage sidecar ahead of time.
-
-Thin wrapper around `giant.pipeline.run_setup_stage` so a dataset's vocab
-maps, event-id split index, and normalizer stats can be warmed once — e.g.
-right after `dwarf convert`, or before kicking off a `dwarf hparam-scan`
-sweep — without needing to also start training. See giant/data/setup_cache.py
-for the sidecar itself.
-"""
-
-from pathlib import Path
-
-from giant.constants import K_MAX
-from giant.pipeline import run_setup_stage
-
-
-def run_warm_setup_cache(
-    data: str,
-    val_fraction: float = 0.1,
-    seed: int = 0,
-    particle_conditioning: str = "physical",
-    material_conditioning: str = "physical",
-    router_enabled: bool = False,
-    router_type: str = "energy",
-    n_experts: int = 4,
-    rebuild: bool = False,
-    echo=print,
-) -> None:
-    """Populate (or refresh) the setup cache sidecar for `data`.
-
-    `val_fraction`/`seed`/`particle_conditioning`/`material_conditioning`
-    select the normalizer cache entry
-    (`giant.data.setup_cache.normalizer_key`) — pass the same values a later
-    `giant train` invocation will use so it hits this warmed entry. The two
-    conditioning axes are independent (docs/v0.3.0-design.md §3.1) and may
-    differ. `router_enabled`/`router_type`/`n_experts` only matter for
-    `router_type == "process"` (warms that `n_experts`'s process map); the
-    energy-router quantile summary is always collected regardless, so a
-    later `--router-type energy` run never needs to rescan just to seed
-    centers.
-    """
-    router_cfg = {
-        "enabled": router_enabled,
-        "type": router_type,
-        "n_experts": n_experts,
-    }
-    # A minimal v0.3 cfg — only the keys run_setup_stage actually reads
-    # (conditioning.{particle,material}.type, stage{1,2}_model.router). This
-    # CLI only ever configures one router (matching today's single
-    # --router-type flag), so it's placed on stage1_model; stage2_model's
-    # stays disabled.
-    cfg = {
-        "conditioning": {
-            "particle": {"type": particle_conditioning},
-            "material": {"type": material_conditioning},
-        },
-        "stage1_model": {"router": router_cfg},
-        "stage2_model": {"router": {"enabled": False}, "k_max": K_MAX},
-    }
-    run_setup_stage(
-        Path(data),
-        val_fraction=val_fraction,
-        seed=seed,
-        cfg=cfg,
-        cache_setup=True,
-        rebuild_setup_cache=rebuild,
-        echo=echo,
-    )
-    echo("setup cache warmed.")
@@ -1,13 +1,13 @@
 """Frozen snapshot of `giant/model/network.py` as it stood at the v0.3.0
-"step 1" commit (eb6dd27), i.e. the last commit before the step-2 §5
-decomposition (see `docs/v0.3.0-design.md`).
+"step 1" commit (eb6dd27), i.e. the last commit before the step-2
+composable-parts decomposition.

 This is a deliberate verbatim copy, not an import of the live module — the
 whole point is that this file's classes keep behaving exactly as v0.2 did
 even after `giant/model/network.py` itself is rewritten, so
 `tests/test_migration_v02_v03.py` has a stable "old" side to diff the new
-`build_models`/`Stage1Model`/`Stage2OneShot` against (design doc §4.3's
-bit-identical acceptance test). Do not edit this file to track future
+`build_models`/`Stage1Model`/`Stage2OneShot` against (the bit-identical
+acceptance test). Do not edit this file to track future
 `network.py` changes — it exists specifically to stop tracking them.
 """

@@ -37,11 +37,7 @@ class SinusoidalEmbedding(nn.Module):
        super().__init__()
        assert dim % 2 == 0, "dim must be even"
        half = dim // 2
-        freqs = torch.exp(
-            -math.log(10000)
-            * torch.arange(half, dtype=torch.float32)
-            / max(half - 1, 1)
-        )
+        freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32) / max(half - 1, 1))
        self.register_buffer("freqs", freqs)

    def forward(self, t: torch.Tensor) -> torch.Tensor:
@@ -107,9 +103,7 @@ class ConditionEncoder(nn.Module):
            pdg_e = self.pdg_emb(cond_cat[:, 0])
            mat_e = self.mat_emb(cond_cat[:, 1])
        else:
-            particle_phys = cond_cont[
-                :, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM
-            ]
+            particle_phys = cond_cont[:, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM]
            material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :]
            pdg_e = self.particle_mlp(particle_phys)
            mat_e = self.material_mlp(material_phys)
@@ -168,12 +162,7 @@ class DenoisingMLP(nn.Module):
        )
        merged_cond_dim = time_dim + cond_out_dim
        self.input_proj = nn.Linear(x_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) for _ in range(n_blocks)])
        self.out_proj = nn.Linear(hidden_dim, x_dim)
        # Predicts n_sec as classification over {0, 1, ..., k_max}.
        # Applied to the condition encoding (not the diffused latent).
@@ -286,12 +275,7 @@ class SecondaryDecoder(nn.Module):
        )
        merged_cond_dim = time_dim + cond_out_dim
        self.input_proj = nn.Linear(sec_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) for _ in range(n_blocks)])
        self.out_proj = nn.Linear(hidden_dim, sec_dim)

    def forward(
@@ -345,12 +329,7 @@ class WGANGenerator(nn.Module):
            conditioning=conditioning,
        )
        self.input_proj = nn.Linear(noise_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)])
        self.out_proj = nn.Linear(hidden_dim, x_dim)
        self.n_sec_head = nn.Sequential(
            nn.Linear(cond_out_dim, hidden_dim // 2),
@@ -410,12 +389,7 @@ class Critic(nn.Module):
            conditioning=conditioning,
        )
        self.input_proj = nn.Linear(x_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)])
        self.out_norm = nn.LayerNorm(hidden_dim)
        self.out_proj = nn.Linear(hidden_dim, 1)

@@ -466,12 +440,7 @@ class WGANSecondaryGenerator(nn.Module):
            conditioning=conditioning,
        )
        self.input_proj = nn.Linear(noise_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)])
        self.out_proj = nn.Linear(hidden_dim, sec_dim)

    def forward(
@@ -515,12 +484,7 @@ class SecondaryCritic(nn.Module):
            conditioning=conditioning,
        )
        self.input_proj = nn.Linear(sec_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)])
        self.out_norm = nn.LayerNorm(hidden_dim)
        self.out_proj = nn.Linear(hidden_dim, 1)

@@ -550,9 +514,7 @@ class Router(nn.Module):
    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        raise NotImplementedError

-    def combine_weights(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
-    ) -> torch.Tensor:
+    def combine_weights(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        probs = self.gate(cond_cont, cond_cat)
        if not (self.gumbel and self.training):
            return probs
@@ -562,26 +524,18 @@ class Router(nn.Module):
    def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        return self.gate(cond_cont, cond_cat).argmax(dim=-1)

-    def balance_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
-    ) -> torch.Tensor:
+    def balance_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        importance = self.gate(cond_cont, cond_cat).sum(dim=0)  # (n_experts,)
        return (importance.std() / (importance.mean() + 1e-8)) ** 2

-    def classify_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
-    ) -> torch.Tensor:
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        return torch.zeros((), device=cond_cont.device)

-    def entropy_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
-    ) -> torch.Tensor:
+    def entropy_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        norm_entropy, _ = self.gate_stats(cond_cont, cond_cat)
        return norm_entropy

-    def gate_stats(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    def gate_stats(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        gate = self.gate(cond_cont, cond_cat)  # (B, n_experts)
        row_entropy = -(gate * (gate + 1e-8).log()).sum(dim=-1)  # (B,)
        norm_entropy = row_entropy.mean() / math.log(self.n_experts)
@@ -602,9 +556,7 @@ def register_router(name: str):

 def build_router(name: str, n_experts: int, **kwargs) -> Router:
    if name not in ROUTER_REGISTRY:
-        raise ValueError(
-            f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}"
-        )
+        raise ValueError(f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}")
    cls = ROUTER_REGISTRY[name]
    accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"}
    filtered = {k: v for k, v in kwargs.items() if k in accepted}
@@ -644,8 +596,7 @@ class EnergyRouter(Router):
        if learn_width or learn_temperature:
            if not (width_min_ratio < 1.0 < width_max_ratio):
                raise ValueError(
-                    f"width_min_ratio ({width_min_ratio}) and width_max_ratio "
-                    f"({width_max_ratio}) must bracket 1.0"
+                    f"width_min_ratio ({width_min_ratio}) and width_max_ratio ({width_max_ratio}) must bracket 1.0"
                )
            self._width_lo = width_min_ratio * temperature
            self._width_hi = width_max_ratio * temperature
@@ -658,10 +609,7 @@ class EnergyRouter(Router):
            centers = torch.linspace(-2.0, 2.0, n_experts)
        else:
            if len(centers_init) != n_experts:
-                raise ValueError(
-                    f"centers_init has {len(centers_init)} values, "
-                    f"expected n_experts={n_experts}"
-                )
+                raise ValueError(f"centers_init has {len(centers_init)} values, expected n_experts={n_experts}")
            centers = torch.tensor(list(centers_init), dtype=torch.float32)
        if learn_centers:
            self.centers = nn.Parameter(centers)
@@ -702,9 +650,7 @@ class PdgRouter(Router):

    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        e = self.pdg_emb(cond_cat[:, 0])  # (B, emb_dim)
-        d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(
-            -1
-        )  # (B, n_experts)
+        d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(-1)  # (B, n_experts)
        return torch.softmax(-d2 / self.temperature, dim=-1)


@@ -736,9 +682,7 @@ class ProcessRouter(Router):
    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        return torch.softmax(self.logits(cond_cont, cond_cat), dim=-1)

-    def classify_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
-    ) -> torch.Tensor:
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        return F.cross_entropy(self.logits(cond_cont, cond_cat), labels)


@@ -756,14 +700,10 @@ class ComposedRouter(Router):
        joint = self.routers[0].gate(cond_cont, cond_cat)  # (B, n_0)
        for router in self.routers[1:]:
            g = router.gate(cond_cont, cond_cat)  # (B, n_i)
-            joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(
-                1
-            )  # (B, prod so far)
+            joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(1)  # (B, prod so far)
        return joint

-    def classify_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
-    ) -> torch.Tensor:
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        total = torch.zeros((), device=cond_cont.device)
        for router in self.routers:
            total = total + router.classify_loss(cond_cont, cond_cat, labels)
@@ -796,12 +736,7 @@ class ExpertTrunk(nn.Module):
    ) -> None:
        super().__init__()
        self.input_proj = nn.Linear(in_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) for _ in range(n_blocks)])
        self.out_proj = nn.Linear(hidden_dim, in_dim)

    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
@@ -891,9 +826,7 @@ class RoutedDenoisingMLP(nn.Module):
        t_emb = self.time_emb(t)
        c_emb = self.cond_enc(cond_cont, cond_cat)
        cond = torch.cat([t_emb, c_emb], dim=-1)
-        return _route_forward(
-            self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
-        )
+        return _route_forward(self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training)

    def predict_n_sec(
        self,
@@ -957,9 +890,7 @@ class RoutedSecondaryDecoder(nn.Module):
        t_emb = self.time_emb(t)
        c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out)
        cond = torch.cat([t_emb, c_emb], dim=-1)
-        return _route_forward(
-            self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
-        )
+        return _route_forward(self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training)


 _STAGE1_MODEL_KEYS = {
@@ -1006,9 +937,7 @@ def _parse_composed_axes(router_cfg: dict) -> list[dict]:
 _VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process")


-def _check_router_conditioning_compat(
-    router_types: list[str], conditioning: str
-) -> None:
+def _check_router_conditioning_compat(router_types: list[str], conditioning: str) -> None:
    bad = sorted(set(router_types) & set(_VOCAB_SCOPED_ROUTER_TYPES))
    if bad and conditioning == "physical":
        raise ValueError(
@@ -1019,9 +948,7 @@ def _check_router_conditioning_compat(
        )


-def _build_router_from_cfg(
-    router_cfg: dict, pdg_vocab: int, mat_vocab: int, conditioning: str = "embedding"
-) -> Router:
+def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int, conditioning: str = "embedding") -> Router:
    shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
    if router_cfg["type"] == "composed":
        axes = _parse_composed_axes(router_cfg)
@@ -1030,9 +957,7 @@ def _build_router_from_cfg(
        router.gumbel = bool(router_cfg.get("gumbel", False))
        return router
    _check_router_conditioning_compat([router_cfg["type"]], conditioning)
-    router_kwargs = {
-        k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")
-    }
+    router_kwargs = {k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")}
    router_kwargs.setdefault("pdg_vocab", pdg_vocab)
    router_kwargs.setdefault("mat_vocab", mat_vocab)
    router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
@@ -1042,15 +967,9 @@ def _build_router_from_cfg(

 def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
    if model_config.get("mode") == "wgan":
-        stage1 = WGANGenerator(
-            **{k: v for k, v in model_config.items() if k in _WGAN_GENERATOR_MODEL_KEYS}
-        )
+        stage1 = WGANGenerator(**{k: v for k, v in model_config.items() if k in _WGAN_GENERATOR_MODEL_KEYS})
        sec_decoder = WGANSecondaryGenerator(
-            **{
-                k: v
-                for k, v in model_config.items()
-                if k in _WGAN_SEC_GENERATOR_MODEL_KEYS
-            }
+            **{k: v for k, v in model_config.items() if k in _WGAN_SEC_GENERATOR_MODEL_KEYS}
        )
        return stage1, sec_decoder

@@ -1061,44 +980,30 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
        shared = dict(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
-            expert_hidden_dim=model_config.get("expert_hidden_dim")
-            or model_config.get("hidden_dim", 128),
-            expert_n_blocks=model_config.get("expert_n_blocks")
-            or model_config.get("n_blocks", 3),
+            expert_hidden_dim=model_config.get("expert_hidden_dim") or model_config.get("hidden_dim", 128),
+            expert_n_blocks=model_config.get("expert_n_blocks") or model_config.get("n_blocks", 3),
            emb_dim=model_config.get("emb_dim", EMB_DIM),
            dropout=model_config.get("dropout", 0.1),
            conditioning=model_config.get("conditioning", "embedding"),
        )
        conditioning = shared["conditioning"]
        stage1 = RoutedDenoisingMLP(
-            router=_build_router_from_cfg(
-                router_cfg, pdg_vocab, mat_vocab, conditioning
-            ),
+            router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, conditioning),
            k_max=model_config.get("k_max", K_MAX),
            **shared,
        )
        sec_decoder = RoutedSecondaryDecoder(
-            router=_build_router_from_cfg(
-                router_cfg, pdg_vocab, mat_vocab, conditioning
-            ),
+            router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, conditioning),
            **shared,
        )
        return stage1, sec_decoder

-    stage1 = DenoisingMLP(
-        **{k: v for k, v in model_config.items() if k in _STAGE1_MODEL_KEYS}
-    )
-    sec_decoder = SecondaryDecoder(
-        **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
-    )
+    stage1 = DenoisingMLP(**{k: v for k, v in model_config.items() if k in _STAGE1_MODEL_KEYS})
+    sec_decoder = SecondaryDecoder(**{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS})
    return stage1, sec_decoder


 def build_critics(model_config: dict) -> tuple[nn.Module, nn.Module]:
-    critic = Critic(
-        **{k: v for k, v in model_config.items() if k in _CRITIC_MODEL_KEYS}
-    )
-    sec_critic = SecondaryCritic(
-        **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
-    )
+    critic = Critic(**{k: v for k, v in model_config.items() if k in _CRITIC_MODEL_KEYS})
+    sec_critic = SecondaryCritic(**{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS})
    return critic, sec_critic
@@ -0,0 +1,136 @@
+"""Tests for giant/training/amp.py (gitea #47)."""
+
+import tempfile
+from pathlib import Path
+
+import pytest
+import torch
+
+from giant.model.routers import EnergyRouter
+from giant.model.wgan import gradient_penalty
+from giant.training.amp import resolve_autocast
+from giant.training.stage2_inputs import _remaining_energy_fraction
+from test_train import _base_cfg, _run_train
+
+# ---------------------------------------------------------------------------
+# resolve_autocast
+# ---------------------------------------------------------------------------
+
+
+def test_resolve_autocast_fp32_is_disabled():
+    device_type, dtype, enabled = resolve_autocast("fp32", torch.device("cpu"))
+    assert device_type == "cpu"
+    assert dtype is torch.float32
+    assert enabled is False
+
+
+def test_resolve_autocast_bf16_on_cpu_is_enabled():
+    """CPU bf16 autocast is what lets the mixed-precision path be tested
+    without a GPU (torch 2.3 supports it)."""
+    device_type, dtype, enabled = resolve_autocast("bf16", torch.device("cpu"))
+    assert device_type == "cpu"
+    assert dtype is torch.bfloat16
+    assert enabled is True
+
+
+def test_resolve_autocast_bf16_on_unsupported_cuda_raises(monkeypatch):
+    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: False)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (7, 0))
+    monkeypatch.setattr(torch.cuda, "get_device_name", lambda device=None: "Tesla V100")
+    with pytest.raises(ValueError, match="bf16"):
+        resolve_autocast("bf16", torch.device("cuda"))
+
+
+def test_resolve_autocast_bf16_on_mps_raises():
+    with pytest.raises(ValueError, match="bf16"):
+        resolve_autocast("bf16", torch.device("mps"))
+
+
+def test_resolve_autocast_unknown_precision_raises():
+    with pytest.raises(ValueError, match="fp32.*bf16"):
+        resolve_autocast("fp16", torch.device("cpu"))
+
+
+# ---------------------------------------------------------------------------
+# End-to-end: train() under bf16 on CPU
+# ---------------------------------------------------------------------------
+
+
+def test_train_end_to_end_bf16_cpu_completes_and_stores_fp32_params():
+    """Reuses tests/test_train.py's synthetic-batch harness — train() itself
+    is device-agnostic, and CPU bf16 autocast is real (not mocked) in torch
+    2.3, so this is a genuine exercise of the autocast region added to
+    FlowDDPMStageTrainer.step/WGANStageTrainer.step, not just a config
+    passthrough check.
+
+    Also asserts the checkpoint's stored parameters are fp32: autocast only
+    changes the dtype of intermediate activations, never the model's own
+    stored weights — a regression here would mean something accidentally
+    cast the model itself (e.g. `model.to(dtype=torch.bfloat16)`) rather than
+    using autocast."""
+    cfg = _base_cfg()
+    cfg["train"]["precision"] = "bf16"
+    with tempfile.TemporaryDirectory() as tmp:
+        out_dir = Path(tmp) / "run"
+        _run_train(cfg, out_dir)
+        assert (out_dir / "last.pt").exists()
+        assert (out_dir / "metrics.csv").exists()
+        ckpt = torch.load(out_dir / "last.pt", weights_only=False)
+        for stage_key in ("model", "sec_decoder"):
+            if stage_key not in ckpt:
+                continue
+            for name, tensor in ckpt[stage_key].items():
+                if tensor.is_floating_point():
+                    assert tensor.dtype == torch.float32, f"{stage_key}.{name} is {tensor.dtype}, expected fp32"
+
+
+@pytest.mark.parametrize("generator", ["wgan", "flow"])
+def test_train_end_to_end_bf16_cpu_stage2_generators(generator):
+    """bf16 covers both trainer subclasses (FlowDDPMStageTrainer and
+    WGANStageTrainer) — the wgan default in _base_cfg exercises the
+    generator-forward/critic-scoring autocast region added to
+    WGANStageTrainer.step, and flow exercises the plain _compute wrap."""
+    cfg = _base_cfg()
+    cfg["train"]["precision"] = "bf16"
+    cfg["stage2_model"]["generator"] = generator
+    with tempfile.TemporaryDirectory() as tmp:
+        _run_train(cfg, Path(tmp) / "run")
+
+
+# ---------------------------------------------------------------------------
+# fp32 guards: correct in fp32, quietly degrade in bf16 — stay fp32 even
+# under an active bf16 autocast region.
+# ---------------------------------------------------------------------------
+
+
+def test_remaining_energy_fraction_stays_fp32_under_bf16_autocast():
+    fraction = torch.rand(4, 5).to(torch.bfloat16)
+    with torch.autocast("cpu", dtype=torch.bfloat16, enabled=True):
+        out = _remaining_energy_fraction(fraction)
+    assert out.dtype == torch.float32
+
+
+def test_gradient_penalty_stays_fp32_under_bf16_autocast():
+    critic = torch.nn.Linear(6, 1)
+
+    def critic_fn(x):
+        return critic(x)
+
+    real = torch.randn(4, 6)
+    fake = torch.randn(4, 6)
+    with torch.autocast("cpu", dtype=torch.bfloat16, enabled=True):
+        gp = gradient_penalty(critic_fn, real, fake)
+    assert gp.dtype == torch.float32
+
+
+def test_router_balance_and_entropy_loss_stay_fp32_under_bf16_autocast():
+    router = EnergyRouter(n_experts=3)
+    cond_cont = torch.randn(8, 15)
+    cond_cat = torch.zeros(8, 2, dtype=torch.long)
+    with torch.autocast("cpu", dtype=torch.bfloat16, enabled=True):
+        balance = router.balance_loss(cond_cont, cond_cat)
+        entropy = router.entropy_loss(cond_cont, cond_cat)
+        weights = router.combine_weights(cond_cont, cond_cat)
+    assert balance.dtype == torch.float32
+    assert entropy.dtype == torch.float32
+    assert weights.dtype == torch.float32
@@ -10,9 +10,11 @@ from giant.analysis import reduce as R
 from giant.analysis.sources import (
    SYNTHETIC_TERMINATION_REASONS,
    Side,
+    open_side,
    physical_steps,
    secondaries,
 )
+from giant.data.loader import EVENT_ID_FILE_STRIDE


 def _rollout_frame() -> pl.LazyFrame:
@@ -102,6 +104,31 @@ def test_hist1d_overall_and_grouped():
    assert hg[11].sum() == 4


+def test_hist1d_clamps_extreme_values_and_drops_nan():
+    # A rollout can emit a wildly out-of-range step_length (or an inf/NaN); the
+    # fixed-edge binning must clamp rather than overflow the i32 bin cast.
+    lf = pl.DataFrame({"x": [5.0, 1.0725e10, float("inf"), -float("inf"), float("nan"), None]}).lazy()
+    edges = np.linspace(0.0, 50.0, 6)  # width 10
+    h = R.hist1d(lf, pl.col("x"), edges)
+    # 5 -> bin 0; 1e10 and +inf -> top bin; -inf -> bin 0; NaN/null dropped
+    assert h[0].tolist() == [2, 0, 0, 0, 2]
+
+
+def test_profile_partial_clamps_extreme_values_and_drops_nan():
+    lf = pl.DataFrame(
+        {
+            "event_id": [1, 1, 1, 1],
+            "z": [5.0, 1.0725e10, float("nan"), 45.0],
+            "w": [1.0, 2.0, 4.0, 8.0],
+        }
+    ).lazy()
+    edges = np.linspace(0.0, 50.0, 6)
+    ev, mat = R.profile_partial(lf, pl.col("z"), edges, pl.col("w"))
+    assert ev.tolist() == [1]
+    # 1e10 clamps into the top bin alongside 45; the NaN row's weight is dropped
+    assert mat[0].tolist() == [1.0, 0.0, 0.0, 0.0, 10.0]
+
+
 def test_physical_steps_drops_synthetic_rollout_rows_only():
    lf = _rollout_frame()
    phys = physical_steps(lf, Side.rollout).collect()
@@ -132,6 +159,20 @@ def test_secondaries_rollout_vs_reference_align():
    assert t["pdg"].to_list() == [22, 22]


+def test_sec_count_by_event_zero_fills_events_with_no_secondaries():
+    r_phys = physical_steps(_rollout_frame(), Side.rollout)
+    r_sec = secondaries(_rollout_frame(), Side.rollout)
+    ev, n = R.sec_count_by_event(r_phys, r_sec)
+    # event 1 has one secondary track; event 2 has none and must still appear (as 0),
+    # not silently drop out of a plain group_by on the secondaries frame alone.
+    assert dict(zip(ev.tolist(), n.tolist())) == {1: 1, 2: 0}
+
+    t_all = _reference_frame()
+    t_sec = secondaries(t_all, Side.reference)
+    ev, n = R.sec_count_by_event(t_all, t_sec)
+    assert dict(zip(ev.tolist(), n.tolist())) == {1: 1, 2: 1}
+
+
 def test_leakage_fraction():
    frac = R.leakage_fraction(_rollout_frame())
    # event 1: escaped pre_E=30, deposited=90 -> 30/120 = 0.25; event 2: 0
@@ -169,3 +210,35 @@ def test_pdg_and_material_labels():
    assert G.pdg_label(22) == "gamma"
    assert G.pdg_label(999999) == "999999"
    assert G.material_label("G4_PbWO4") == "PbWO4"
+
+
+def _write_shard(path, event_ids, edeps):
+    pl.DataFrame({"event_id": event_ids, "pdg": [11] * len(event_ids), "edep": edeps}).write_parquet(path)
+
+
+def test_open_side_reference_offsets_event_ids_across_shards(tmp_path):
+    # Each shard is a separate Geant4 job whose own event_id numbering restarts
+    # from 0 — a naive multi-shard scan collides on event_id across shards.
+    _write_shard(tmp_path / "a.parquet", [0, 1], [1.0, 2.0])
+    _write_shard(tmp_path / "b.parquet", [0, 1], [3.0, 4.0])
+    df = open_side(tmp_path, Side.reference).sort("event_id").collect()
+    assert df["event_id"].to_list() == [0, 1, EVENT_ID_FILE_STRIDE, EVENT_ID_FILE_STRIDE + 1]
+    assert df["edep"].to_list() == [1.0, 2.0, 3.0, 4.0]
+    assert "__source_path" not in df.columns
+
+
+def test_open_side_reference_single_file_unchanged(tmp_path):
+    _write_shard(tmp_path / "only.parquet", [0, 1], [1.0, 2.0])
+    df = open_side(tmp_path / "only.parquet", Side.reference).sort("event_id").collect()
+    assert df["event_id"].to_list() == [0, 1]
+    assert "__source_path" not in df.columns
+
+
+def test_open_side_reference_manifest(tmp_path):
+    _write_shard(tmp_path / "a.parquet", [0, 1], [1.0, 2.0])
+    _write_shard(tmp_path / "b.parquet", [0, 1], [3.0, 4.0])
+    manifest = tmp_path / "shards.manifest"
+    manifest.write_text("a.parquet\nb.parquet\n")
+    df = open_side(manifest, Side.reference).sort("event_id").collect()
+    assert df["event_id"].to_list() == [0, 1, EVENT_ID_FILE_STRIDE, EVENT_ID_FILE_STRIDE + 1]
+    assert df["edep"].to_list() == [1.0, 2.0, 3.0, 4.0]
@@ -1,8 +1,7 @@
-"""Tests for the rollout-YAML → run-directory flow, compute, and submit."""
+"""Tests for the rollout-YAML(s) → run-directory flow, compute, and merge."""

 from __future__ import annotations

-import sys
 from pathlib import Path

 import pyarrow.parquet as pq
@@ -11,30 +10,31 @@ import yaml

 from giant.analysis import (
    RunMeta,
-    SubmitConfig,
-    catalog_ids,
    compute_one,
    compute_reduced,
    derive_run_dir,
    load_rollout_yaml,
+    load_rollout_yamls,
    merge_one,
    prep,
-    write_submit,
 )
-from giant.analysis.catalog import get_spec
-from giant.analysis.condor import Context
+from giant.analysis.run import Context
 from giant.analysis.reduced import Partial, Reduced
 from giant.constants import PREDICT_COORD_METADATA_KEY, ROLLOUT_COORD_VALUE
 from tests.test_analysis_reduce import _reference_frame, _rollout_frame


+def _write_rollout(path: Path) -> None:
+    tbl = _rollout_frame().collect().to_arrow()
+    tbl = tbl.replace_schema_metadata({PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE})
+    pq.write_table(tbl, path)
+
+
 def _write_inputs(tmp_path: Path) -> Path:
    """Materialize rollout+reference parquet and a rollout YAML; return the YAML path."""
    rollout = tmp_path / "rollout.parquet"
    reference = tmp_path / "reference.parquet"
-    tbl = _rollout_frame().collect().to_arrow()
-    tbl = tbl.replace_schema_metadata({PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE})
-    pq.write_table(tbl, rollout)
+    _write_rollout(rollout)
    _reference_frame().collect().write_parquet(reference)

    yaml_path = tmp_path / "run.yaml"
@@ -54,22 +54,40 @@ def _write_inputs(tmp_path: Path) -> Path:
    return yaml_path


-def _fake_venv(repo_dir: Path) -> None:
-    """Stand in for a `uv sync`'d venv: write_submit checks `.venv/bin/giant` exists."""
-    giant = repo_dir / ".venv" / "bin" / "giant"
-    giant.parent.mkdir(parents=True, exist_ok=True)
-    giant.write_text("#!/bin/bash\n")
-    giant.chmod(0o755)
+def _write_two_inputs(tmp_path: Path) -> tuple[Path, Path]:
+    """Two rollout YAMLs (distinct output files) sharing one reference file."""
+    reference = tmp_path / "reference.parquet"
+    _reference_frame().collect().write_parquet(reference)
+
+    paths = []
+    for tag, pred_id in (("a", "aaaa1111ef"), ("b", "bbbb2222ef")):
+        rollout = tmp_path / f"rollout_{tag}.parquet"
+        _write_rollout(rollout)
+        yaml_path = tmp_path / f"run_{tag}.yaml"
+        yaml_path.write_text(
+            yaml.safe_dump(
+                {
+                    "prediction_id": pred_id,
+                    "output": str(rollout),
+                    "dataset": str(reference),
+                    "checkpoint": f"/ckpt/{tag}.pt",
+                    "kind": "rollout",
+                    "energy_cutoff": 0.1,
+                    "steps": 10,
+                }
+            )
+        )
+        paths.append(yaml_path)
+    return paths[0], paths[1]


-def _prep(
-    rollout_yaml: Path, run_dir: str | Path | None = None, chunks: int = 1
-) -> Path:
+def _prep(rollout_yamls, run_dir: str | Path | None = None, chunks: int = 1, labels=None) -> Path:
    """``prep`` with small test-sized context bins/sampling."""
    return prep(
-        rollout_yaml,
+        rollout_yamls,
        run_dir,
        n_chunks=chunks,
+        labels=labels,
        n_energy_bins=2,
        n_marginal_bins=8,
        top_k_pdg=3,
@@ -84,39 +102,108 @@ def test_load_rollout_yaml_requires_paths(tmp_path: Path):
        load_rollout_yaml(bad)


+def test_load_rollout_yamls_single_defaults_to_rollout_name(tmp_path: Path):
+    yaml_path = _write_inputs(tmp_path)
+    loaded, reference = load_rollout_yamls([yaml_path])
+    assert [lr.name for lr in loaded] == ["rollout"]
+    assert reference.endswith("reference.parquet")
+
+
+def test_load_rollout_yamls_multi_defaults_to_stem(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    loaded, _ = load_rollout_yamls([a, b])
+    assert [lr.name for lr in loaded] == ["run_a", "run_b"]
+
+
+def test_load_rollout_yamls_explicit_labels(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    loaded, _ = load_rollout_yamls([a, b], labels=["flow", "wgan"])
+    assert [lr.name for lr in loaded] == ["flow", "wgan"]
+
+
+def test_load_rollout_yamls_label_count_mismatch(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    with pytest.raises(ValueError, match="--label"):
+        load_rollout_yamls([a, b], labels=["only-one"])
+
+
+def test_load_rollout_yamls_rejects_duplicate_names(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    with pytest.raises(ValueError, match="collide"):
+        load_rollout_yamls([a, b], labels=["same", "same"])
+
+
+def test_load_rollout_yamls_rejects_mismatched_reference(tmp_path: Path):
+    a, _ = _write_two_inputs(tmp_path)
+    other_ref = tmp_path / "other_reference.parquet"
+    _reference_frame().collect().write_parquet(other_ref)
+    c = tmp_path / "run_c.yaml"
+    c.write_text(
+        yaml.safe_dump(
+            {"prediction_id": "cccc3333ef", "output": str(tmp_path / "rollout_c.parquet"), "dataset": str(other_ref)}
+        )
+    )
+    _write_rollout(tmp_path / "rollout_c.parquet")
+    with pytest.raises(ValueError, match="same reference"):
+        load_rollout_yamls([a, c])
+
+
 def test_derive_run_dir_next_to_rollout():
    y = {"output": "/data/roll.parquet", "prediction_id": "abcd1234ef", "dataset": "d"}
-    assert derive_run_dir(y) == Path("/data/analysis_abcd1234")
-    assert derive_run_dir(y, "/somewhere") == Path("/somewhere")
+    assert derive_run_dir([y]) == Path("/data/analysis_abcd1234")
+    assert derive_run_dir([y], "/somewhere") == Path("/somewhere")


 def test_derive_run_dir_default_base():
    y = {"output": "/data/roll.parquet", "prediction_id": "abcd1234ef", "dataset": "d"}
-    assert derive_run_dir(y, default_base="/work/lbogner/giant2/analysis_runs") == Path(
+    assert derive_run_dir([y], default_base="/work/lbogner/giant2/analysis_runs") == Path(
        "/work/lbogner/giant2/analysis_runs/analysis_abcd1234"
    )
    # an explicit run_dir still wins over default_base
-    assert derive_run_dir(y, "/somewhere", default_base="/other") == Path("/somewhere")
+    assert derive_run_dir([y], "/somewhere", default_base="/other") == Path("/somewhere")
+
+
+def test_derive_run_dir_multi_rollout_joins_tags():
+    ys = [{"output": f"/data/roll_{i}.parquet", "prediction_id": f"tag{i}xxxx", "dataset": "d"} for i in range(2)]
+    assert derive_run_dir(ys, default_base="/base") == Path("/base/analysis_tag0xxxx-tag1xxxx")
+
+
+def test_derive_run_dir_many_rollouts_truncates_with_plus_count():
+    ys = [{"output": f"/data/roll_{i}.parquet", "prediction_id": f"tag{i}xxxx", "dataset": "d"} for i in range(5)]
+    run_dir = derive_run_dir(ys, default_base="/base")
+    assert run_dir == Path("/base/analysis_tag0xxxx-tag1xxxx-tag2xxxx-plus2")


 def test_prep_lays_out_run_dir(tmp_path: Path):
    yaml_path = _write_inputs(tmp_path)
-    run_dir = _prep(yaml_path)
+    run_dir = _prep([yaml_path])
    assert run_dir == tmp_path / "analysis_abcd1234"
    assert (run_dir / "shared.json").exists()
    ctx = Context.load(run_dir / "shared.json")
    assert set(ctx.var_ranges) == {"step_length", "edep", "delta_e", "post_E"}
    meta = RunMeta.load(run_dir / "run_meta.json")
    assert meta.reference.endswith("reference.parquet")
-    assert meta.plot_meta["checkpoint"] == "/ckpt/best.pt"
+    assert [ro["name"] for ro in meta.rollouts] == ["rollout"]
+    assert meta.rollouts[0]["plot_meta"]["checkpoint"] == "/ckpt/best.pt"
    assert "best.pt" in meta.title
    assert meta.n_chunks == 1
    assert meta.rows_per_chunk == [meta.total_rows]  # single chunk holds everything
    assert meta.total_rows == 8  # 5 rollout rows + 3 reference rows


+def test_prep_multi_rollout_lays_out_run_dir(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    run_dir = _prep([a, b], labels=["flow", "wgan"])
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    assert [ro["name"] for ro in meta.rollouts] == ["flow", "wgan"]
+    assert meta.rollouts[0]["plot_meta"]["checkpoint"] == "/ckpt/a.pt"
+    assert meta.rollouts[1]["plot_meta"]["checkpoint"] == "/ckpt/b.pt"
+    # 5 rows from each rollout + 3 from the shared reference
+    assert meta.total_rows == 13
+
+
 def test_prep_splits_rows_per_chunk(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path), chunks=2)
+    run_dir = _prep([_write_inputs(tmp_path)], chunks=2)
    meta = RunMeta.load(run_dir / "run_meta.json")
    assert len(meta.rows_per_chunk) == 2
    assert sum(meta.rows_per_chunk) == meta.total_rows == 8
@@ -127,7 +214,7 @@ def test_reprep_clears_stale_partials_from_a_different_chunk_count(tmp_path: Pat
    partials on disk for merge_one to silently merge against the new
    context (they'd be keyed/sized for the old n_chunks)."""
    yaml_path = _write_inputs(tmp_path)
-    run_dir = _prep(yaml_path, chunks=2)
+    run_dir = _prep([yaml_path], chunks=2)
    compute_one("marginal_edep", run_dir, chunk_index=0)
    compute_one("marginal_edep", run_dir, chunk_index=1)
    stale = run_dir / "reduced_partial" / "marginal_edep__0.json"
@@ -135,7 +222,7 @@ def test_reprep_clears_stale_partials_from_a_different_chunk_count(tmp_path: Pat
    (run_dir / "reduced").mkdir(exist_ok=True)
    (run_dir / "reduced" / "marginal_edep.json").write_text("{}")

-    _prep(yaml_path, run_dir, chunks=1)
+    _prep([yaml_path], run_dir, chunks=1)

    assert not stale.exists()
    assert not (run_dir / "reduced" / "marginal_edep.json").exists()
@@ -143,20 +230,22 @@ def test_reprep_clears_stale_partials_from_a_different_chunk_count(tmp_path: Pat


 def test_compute_one_from_run_dir(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path))
+    run_dir = _prep([_write_inputs(tmp_path)])
    out = compute_one("marginal_edep", run_dir)
    assert out == run_dir / "reduced_partial" / "marginal_edep__0.json"
    partial = Partial.load(out)
    assert partial.id == "marginal_edep" and partial.chunk == 0
    assert "r" in partial.data and "t" in partial.data
+    assert list(partial.data["r"]) == ["rollout"]


 def test_compute_reduced_explicit_paths(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path))
+    run_dir = _prep([_write_inputs(tmp_path)])
    meta = RunMeta.load(run_dir / "run_meta.json")
+    rollouts = [{"name": ro["name"], "path": ro["path"]} for ro in meta.rollouts]
    out = compute_reduced(
        "marginal_step_length",
-        meta.rollout,
+        rollouts,
        meta.reference,
        run_dir / "shared.json",
        tmp_path / "r.json",
@@ -165,17 +254,17 @@ def test_compute_reduced_explicit_paths(tmp_path: Path):


 def test_merge_one_produces_reduced(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path))
+    run_dir = _prep([_write_inputs(tmp_path)])
    compute_one("marginal_edep", run_dir)
    out = merge_one("marginal_edep", run_dir)
    assert out == run_dir / "reduced" / "marginal_edep.json"
    reduced = Reduced.load(out)
    assert reduced.id == "marginal_edep"
-    assert len(reduced.payload["rollout"]) == len(reduced.payload["edges"]) - 1
+    assert len(reduced.payload["series"]["rollout"]) == len(reduced.payload["edges"]) - 1


 def test_merge_one_fails_loudly_on_missing_chunk(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path), chunks=2)
+    run_dir = _prep([_write_inputs(tmp_path)], chunks=2)
    compute_one("marginal_edep", run_dir, chunk_index=0)  # chunk 1 never computed
    with pytest.raises(FileNotFoundError, match="missing chunk"):
        merge_one("marginal_edep", run_dir)
@@ -184,11 +273,11 @@ def test_merge_one_fails_loudly_on_missing_chunk(tmp_path: Path):
 def test_chunked_compute_and_merge_matches_unchunked(tmp_path: Path):
    (tmp_path / "a").mkdir()
    (tmp_path / "b").mkdir()
-    unchunked_dir = _prep(_write_inputs(tmp_path / "a"))
+    unchunked_dir = _prep([_write_inputs(tmp_path / "a")])
    compute_one("marginal_step_length", unchunked_dir)
    unchunked = Reduced.load(merge_one("marginal_step_length", unchunked_dir))

-    chunked_dir = _prep(_write_inputs(tmp_path / "b"), chunks=2)
+    chunked_dir = _prep([_write_inputs(tmp_path / "b")], chunks=2)
    for k in range(2):
        compute_one("marginal_step_length", chunked_dir, chunk_index=k)
    chunked = Reduced.load(merge_one("marginal_step_length", chunked_dir))
@@ -196,89 +285,21 @@ def test_chunked_compute_and_merge_matches_unchunked(tmp_path: Path):
    assert chunked.payload == unchunked.payload


+def test_two_rollout_compute_and_merge_produces_both_series(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    run_dir = _prep([a, b], labels=["flow", "wgan"])
+    compute_one("marginal_edep", run_dir)
+    reduced = Reduced.load(merge_one("marginal_edep", run_dir))
+    assert list(reduced.payload["series"]) == ["flow", "wgan"]
+    assert "reference" in reduced.payload
+
+
 def test_compute_reduced_rejects_out_of_range_chunk(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path))  # n_chunks=1 (default)
+    run_dir = _prep([_write_inputs(tmp_path)])  # n_chunks=1 (default)
    with pytest.raises(ValueError, match="out of range"):
        compute_one("marginal_edep", run_dir, chunk_index=1)


-def test_write_submit_description(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path))
-    _fake_venv(tmp_path)
-    cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path)
-    txt = write_submit(cfg).read_text()
-    assert "universe = docker" in txt
-    assert "docker_image = cverstege/alma9-gridjob" in txt
-    assert "requirements = TARGET.ProvidesETPResources" in txt
-    assert "accounting_group = cms" in txt
-    assert "+RequestWalltime = $(walltime)" in txt
-    assert "queue plotid,chunk,walltime from" in txt
-    jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()]
-    assert [i for i, _, _ in jobs] == catalog_ids()
-    assert all(k == "0" for _, k, _ in jobs)  # n_chunks=1 default
-    assert all(int(w) > 0 for _, _, w in jobs)
-    wrapper = run_dir / "run_compute.sh"
-    assert wrapper.exists() and (wrapper.stat().st_mode & 0o111)
-    body = wrapper.read_text()
-    assert "giant analyze compute-one --id" in body
-    assert "--chunk" in body and "--run-dir" in body
-
-
-def test_write_submit_requires_synced_venv(
-    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
-):
-    run_dir = _prep(_write_inputs(tmp_path))
-    cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path)
-    # No `giant` next to the (fake) active interpreter, so this falls through
-    # to repo_dir/.venv/bin/giant, which _write_inputs/_prep also didn't create.
-    monkeypatch.setattr(
-        sys, "executable", str(tmp_path / "not-a-venv" / "bin" / "python")
-    )
-    with pytest.raises(FileNotFoundError, match="uv sync"):
-        write_submit(cfg)
-
-
-def test_write_submit_remote_flag(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path))
-    _fake_venv(tmp_path)
-    cfg = SubmitConfig(
-        run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, remote=True
-    )
-    txt = write_submit(cfg).read_text()
-    assert "+RemoteJob = True" in txt
-    assert "ProvidesETPResources" not in txt
-
-
-def test_write_submit_chunks_respect_chunkable(tmp_path: Path):
-    assert get_spec("router_gating").chunkable is False
-    run_dir = _prep(_write_inputs(tmp_path), chunks=4)
-    _fake_venv(tmp_path)
-    cfg = SubmitConfig(
-        run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4
-    )
-    write_submit(cfg)
-    jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()]
-    counts: dict[str, int] = {}
-    for spec_id, _, _ in jobs:
-        counts[spec_id] = counts.get(spec_id, 0) + 1
-    assert counts["marginal_edep"] == 4
-    assert counts["router_gating"] == 1  # chunkable=False, ignores n_chunks
-
-
-def test_write_submit_rejects_n_chunks_mismatch_with_run_meta(tmp_path: Path):
-    """cfg.n_chunks must match the n_chunks the run_dir was actually prepped
-    with — RunMeta.rows_per_chunk is sized to the prepped value, so a
-    mismatch would otherwise surface as a confusing IndexError deep inside
-    _job_walltimes instead of a clear error here."""
-    run_dir = _prep(_write_inputs(tmp_path), chunks=2)
-    _fake_venv(tmp_path)
-    cfg = SubmitConfig(
-        run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4
-    )
-    with pytest.raises(ValueError, match="n_chunks"):
-        write_submit(cfg)
-
-
 def test_estimate_runtime_s_scales_with_rows_and_margin():
    from giant.analysis import RUNTIME_SAFETY_MARGIN, estimate_runtime_s
    from giant.analysis.runtime_estimate import _FIXED_OVERHEAD_S
@@ -288,25 +309,3 @@ def test_estimate_runtime_s_scales_with_rows_and_margin():
    large = estimate_runtime_s("marginal_edep", 100_000_000)
    assert small >= (1 + RUNTIME_SAFETY_MARGIN) * _FIXED_OVERHEAD_S
    assert large > small  # bigger chunk -> longer estimate
-
-
-def test_write_submit_walltime_grows_with_chunk_rows(tmp_path: Path):
-    """A chunked run's later job walltimes track that chunk's row count."""
-    from giant.analysis.runtime_estimate import estimate_runtime_s
-
-    run_dir = _prep(_write_inputs(tmp_path), chunks=2)
-    meta = RunMeta.load(run_dir / "run_meta.json")
-    _fake_venv(tmp_path)
-    cfg = SubmitConfig(
-        run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=2
-    )
-    write_submit(cfg)
-    jobs = {
-        (i, int(k)): int(w)
-        for i, k, w in (
-            line.split(",") for line in (run_dir / "jobs.txt").read_text().split()
-        )
-    }
-    for chunk in range(2):
-        expected = estimate_runtime_s("marginal_edep", meta.rows_per_chunk[chunk])
-        assert jobs[("marginal_edep", chunk)] == expected
@@ -1,7 +1,7 @@
 import os
 import subprocess

-from scripts import bump_dataset_version
+from giant.tools import bump_dataset_version

 plan_bump_gen = bump_dataset_version.plan_bump_gen
 plan_bump_schema = bump_dataset_version.plan_bump_schema
@@ -22,9 +22,7 @@ def test_git_user_name_returns_none_on_timeout(monkeypatch):


 def test_bump_gen_starts_at_gen1_when_none_exist(tmp_path):
-    dirs, log_line = plan_bump_gen(
-        tmp_path, "steps", "first generation", None, "2026-01-01"
-    )
+    dirs, log_line = plan_bump_gen(tmp_path, "steps", "first generation", None, "2026-01-01")
    assert dirs == [
        tmp_path / "raw" / "steps" / "gen1",
        tmp_path / "processed" / "steps" / "gen1" / "schema1",
@@ -56,9 +54,7 @@ def test_bump_gen_kinds_are_independent(tmp_path):

 def test_bump_schema_starts_at_schema1_for_a_fresh_gen(tmp_path):
    (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
-    dirs, log_line = plan_bump_schema(
-        tmp_path, "steps", "gen1", "added e_sec column", None, "2026-01-01"
-    )
+    dirs, log_line = plan_bump_schema(tmp_path, "steps", "gen1", "added e_sec column", None, "2026-01-01")
    assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema1"]
    assert "`gen1`/`schema1`" in log_line

@@ -66,18 +62,14 @@ def test_bump_schema_starts_at_schema1_for_a_fresh_gen(tmp_path):
 def test_bump_schema_increments_within_its_gen(tmp_path):
    (tmp_path / "processed" / "steps" / "gen1" / "schema1").mkdir(parents=True)
    (tmp_path / "processed" / "steps" / "gen1" / "schema2").mkdir(parents=True)
-    dirs, _ = plan_bump_schema(
-        tmp_path, "steps", "gen1", "next schema", None, "2026-01-01"
-    )
+    dirs, _ = plan_bump_schema(tmp_path, "steps", "gen1", "next schema", None, "2026-01-01")
    assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema3"]


 def test_bump_schema_does_not_see_other_gens_schemas(tmp_path):
    (tmp_path / "processed" / "steps" / "gen1" / "schema5").mkdir(parents=True)
    (tmp_path / "raw" / "steps" / "gen2").mkdir(parents=True)
-    dirs, _ = plan_bump_schema(
-        tmp_path, "steps", "gen2", "fresh schema for gen2", None, "2026-01-01"
-    )
+    dirs, _ = plan_bump_schema(tmp_path, "steps", "gen2", "fresh schema for gen2", None, "2026-01-01")
    assert dirs == [tmp_path / "processed" / "steps" / "gen2" / "schema1"]


@@ -91,9 +83,7 @@ def test_bump_schema_rejects_nonexistent_gen(tmp_path):

 def test_bump_gen_to_specific_tag(tmp_path):
    (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
-    dirs, log_line = plan_bump_gen(
-        tmp_path, "steps", "jump to gen5", None, "2026-01-01", target="gen5"
-    )
+    dirs, log_line = plan_bump_gen(tmp_path, "steps", "jump to gen5", None, "2026-01-01", target="gen5")
    assert dirs[0] == tmp_path / "raw" / "steps" / "gen5"
    assert "`gen5`" in log_line

@@ -125,9 +115,7 @@ def test_bump_schema_to_specific_tag(tmp_path):
 def test_bump_schema_rejects_invalid_to_tag(tmp_path):
    (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
    try:
-        plan_bump_schema(
-            tmp_path, "steps", "gen1", "bad tag", None, "2026-01-01", target="v3"
-        )
+        plan_bump_schema(tmp_path, "steps", "gen1", "bad tag", None, "2026-01-01", target="v3")
        assert False, "expected SystemExit"
    except SystemExit:
        pass
@@ -164,15 +152,7 @@ def _make_parquet(path):


 def test_update_manifest_bumps_to_specified_schema(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
    _make_parquet(parquet)

    manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -194,15 +174,7 @@ def test_update_manifest_auto_detects_highest_schema(tmp_path):
    for schema in ("schema1", "schema2", "schema3"):
        d = tmp_path / "processed" / "steps" / "gen1" / schema / "pbwo4"
        d.mkdir(parents=True)
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema3"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet"
    parquet.touch()

    manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -231,15 +203,7 @@ def test_update_manifest_reports_missing_targets(tmp_path):


 def test_update_manifest_skips_already_at_target(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
    _make_parquet(parquet)

    manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -254,15 +218,7 @@ def test_update_manifest_skips_already_at_target(tmp_path):


 def test_update_manifest_preserves_comments_and_blanks(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
    _make_parquet(parquet)

    manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -278,15 +234,7 @@ def test_update_manifest_preserves_comments_and_blanks(tmp_path):


 def test_update_manifest_bumps_gen(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen2"
-        / "schema1"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen2" / "schema1" / "pbwo4" / "shard-000.parquet"
    _make_parquet(parquet)

    manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -303,15 +251,7 @@ def test_update_manifest_bumps_gen(tmp_path):


 def test_update_manifest_bumps_gen_and_schema(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen2"
-        / "schema3"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen2" / "schema3" / "pbwo4" / "shard-000.parquet"
    _make_parquet(parquet)

    manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -328,15 +268,7 @@ def test_update_manifest_bumps_gen_and_schema(tmp_path):


 def test_apply_update_manifest_writes_file(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
    _make_parquet(parquet)

    manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -358,24 +290,8 @@ def test_apply_update_manifest_writes_file(tmp_path):


 def test_create_manifest_writes_relative_paths(tmp_path):
-    pq1 = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
-    pq2 = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-001.parquet"
-    )
+    pq1 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
+    pq2 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-001.parquet"
    _make_parquet(pq1)
    _make_parquet(pq2)

@@ -419,9 +335,7 @@ def test_run_create_manifest_refuses_to_overwrite_existing_output(tmp_path):
    output.write_text("original contents\n")

    try:
-        bump_dataset_version.run_create_manifest(
-            [str(pq)], execute=True, output=str(output)
-        )
+        bump_dataset_version.run_create_manifest([str(pq)], execute=True, output=str(output))
        assert False, "expected SystemExit"
    except SystemExit:
        pass
@@ -435,9 +349,7 @@ def test_run_create_manifest_force_overwrites_existing_output(tmp_path):
    output.parent.mkdir(parents=True)
    output.write_text("original contents\n")

-    bump_dataset_version.run_create_manifest(
-        [str(pq)], execute=True, output=str(output), force=True
-    )
+    bump_dataset_version.run_create_manifest([str(pq)], execute=True, output=str(output), force=True)
    assert output.read_text() != "original contents\n"


@@ -6,26 +6,50 @@ import numpy as np
 import pytest

 from giant.analysis import build_catalog, catalog_ids, get_spec
-from giant.analysis.catalog import Bundle, PlotSpec
+from giant.analysis.catalog import (
+    Bundle,
+    PlotSpec,
+    _containment_depths,
+    _integer_confusion,
+    _ks_statistic,
+)
 from giant.analysis.context import Context, build_context
+from giant.analysis.sources import RolloutSpec
 from tests.test_analysis_reduce import _reference_frame, _rollout_frame


 def _build_ctx() -> Context:
    r, t = _rollout_frame(), _reference_frame()
    return build_context(
-        r, t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000
+        [RolloutSpec("rollout", r)], t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000
    )


+def _two_rollout_specs() -> list[RolloutSpec]:
+    # Two distinct rollout sources so multi-series merging/finalize code is
+    # exercised even though the underlying frame is the same fixture.
+    return [RolloutSpec("flow", _rollout_frame()), RolloutSpec("wgan", _rollout_frame())]
+
+
@pytest.fixture(scope="module")
 def ctx() -> Context:
    return _build_ctx()


+@pytest.fixture(scope="module")
+def two_ctx() -> Context:
+    t = _reference_frame()
+    return build_context(_two_rollout_specs(), t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000)
+
+
@pytest.fixture(scope="module")
 def bundle(ctx: Context) -> Bundle:
-    return Bundle.open(_rollout_frame(), _reference_frame(), ctx)
+    return Bundle.open([RolloutSpec("rollout", _rollout_frame())], _reference_frame(), ctx)
+
+
+@pytest.fixture(scope="module")
+def two_bundle(two_ctx: Context) -> Bundle:
+    return Bundle.open(_two_rollout_specs(), _reference_frame(), two_ctx)


 def test_catalog_ids_unique_and_nonempty():
@@ -55,41 +79,76 @@ def test_every_spec_computes_valid_reduced(bundle: Bundle):
            "single_hist",
            "router_gating",
            "router_share",
+            "router_specialization",
+            "heatmap",
            "unavailable",
        }
        assert r.title and r.xlabel
-        _validate_payload(r)
+        _validate_payload(r, ["rollout"])


-def _validate_payload(r) -> None:
+def test_every_spec_computes_valid_reduced_with_two_rollouts(two_bundle: Bundle):
+    for spec in build_catalog():
+        r = spec.finalize([spec.compute_partial(two_bundle)], two_bundle.ctx)
+        assert r.id == spec.id
+        _validate_payload(r, ["flow", "wgan"])
+
+
+def _validate_payload(r, names: list[str]) -> None:
    p = r.payload
    if r.kind == "overlay_hist":
        n = len(p["edges"]) - 1
-        assert len(p["rollout"]) == n and len(p["reference"]) == n
+        assert list(p["series"]) == names
+        for v in p["series"].values():
+            assert len(v) == n
+        assert len(p["reference"]) == n
    elif r.kind == "single_hist":
-        assert len(p["rollout"]) == len(p["edges"]) - 1
+        assert list(p["series"]) == names
+        for v in p["series"].values():
+            assert len(v) == len(p["edges"]) - 1
    elif r.kind == "grouped_hist":
        n = len(p["edges"]) - 1
        assert p["groups"], "grouped hist must have at least one group"
        for g in p["groups"].values():
-            assert len(g["rollout"]) == n and len(g["reference"]) == n
+            assert list(g["series"]) == names
+            for v in g["series"].values():
+                assert len(v) == n
+            assert len(g["reference"]) == n
    elif r.kind == "profile":
        n = len(p["edges"]) - 1
-        for k in ("rollout_mean", "rollout_std", "reference_mean", "reference_std"):
-            assert len(p[k]) == n
+        assert list(p["series"]) == names
+        for side in p["series"].values():
+            assert len(side["mean"]) == n and len(side["std"]) == n
+        assert len(p["reference"]["mean"]) == n and len(p["reference"]["std"]) == n
    elif r.kind == "bar":
-        assert len(p["labels"]) == len(p["rollout"]) == len(p["reference"])
+        assert list(p["series"]) == names
+        for v in p["series"].values():
+            assert len(p["labels"]) == len(v)
+        assert len(p["labels"]) == len(p["reference"])
    elif r.kind == "unavailable":
        assert p["note"]
    elif r.kind == "router_gating":
-        for side in ("rollout", "reference"):
-            if side in p:
-                assert len(p[side]["centers"]) == len(p[side]["means"])
-    elif r.kind == "router_share":
-        for cat in p["categories"]:
+        for entry in p["series"].values():
            for side in ("rollout", "reference"):
-                if side in p:
-                    assert cat in p[side]
+                if side in entry:
+                    assert len(entry[side]["centers"]) == len(entry[side]["means"])
+    elif r.kind == "router_share":
+        for entry in p["series"].values():
+            for cat in entry["categories"]:
+                for side in ("rollout", "reference"):
+                    if side in entry:
+                        assert cat in entry[side]
+    elif r.kind == "router_specialization":
+        for entry in p["series"].values():
+            for side in ("rollout", "reference"):
+                if side in entry:
+                    assert len(entry[side]["centers"]) == len(entry[side]["score"])
+    elif r.kind == "heatmap":
+        assert list(p["series"]) == names
+        for mat in p["series"].values():
+            assert len(mat) == len(p["row_labels"])
+            for row in mat:
+                assert len(row) == len(p["col_labels"])


 # ---------------------------------------------------------------------------
@@ -100,7 +159,10 @@ def _validate_payload(r) -> None:
 # sec_count_per_species via pdg-keyed sums), concat-then-finalize with
 # data-dependent edges (event_total_edep), concat-then-mean/std (shower_
 # longitudinal), concat-then-max-edge (leakage_fraction), pdg-keyed sum with a
-# ratio (species_edep_share), and a chunkable=False passthrough (router_gating).
+# ratio (species_edep_share), a chunkable=False passthrough (router_gating),
+# nested sum-merge into a scorecard (marginal_distance_summary), concat-then-
+# event-id-join (n_sec_confusion), and concat-then-per-event-derived-quantity
+# (shower_containment_depth_90, reusing the profile matrix's own merge shape).
 _CHUNK_EQUIVALENCE_IDS = [
    "marginal_edep",
    "species_edep_share",
@@ -109,6 +171,9 @@ _CHUNK_EQUIVALENCE_IDS = [
    "leakage_fraction",
    "sec_count_per_species",
    "router_gating",
+    "marginal_distance_summary",
+    "n_sec_confusion",
+    "shower_containment_depth_90",
 ]


@@ -130,24 +195,81 @@ def _assert_payload_close(a, b, path: str = "payload") -> None:


@pytest.mark.parametrize("spec_id", _CHUNK_EQUIVALENCE_IDS)
-def test_chunked_matches_unchunked(ctx: Context, spec_id: str):
+def test_chunked_matches_unchunked(two_ctx: Context, spec_id: str):
    """A plot computed over N event-disjoint chunks then merged must equal the
    same plot computed in one unchunked pass — the core chunking correctness
-    guarantee (see the analysis-rollout-plots chunking plan)."""
+    guarantee (see the analysis-rollout-plots chunking plan). Exercised with
+    two rollout series so the per-rollout merge path is covered too."""
    spec: PlotSpec = get_spec(spec_id)
-    r, t = _rollout_frame(), _reference_frame()
+    rollouts, t = _two_rollout_specs(), _reference_frame()

-    unchunked_bundle = Bundle.open(r, t, ctx)
-    unchunked = spec.finalize([spec.compute_partial(unchunked_bundle)], ctx)
+    unchunked_bundle = Bundle.open(rollouts, t, two_ctx)
+    unchunked = spec.finalize([spec.compute_partial(unchunked_bundle)], two_ctx)

    # 4 chunks over only 2 distinct event_ids also exercises empty chunks.
    n_chunks = 4 if spec.chunkable else 1
-    parts = [
-        spec.compute_partial(Bundle.open(r, t, ctx, chunk=(k, n_chunks)))
-        for k in range(n_chunks)
-    ]
-    chunked = spec.finalize(parts, ctx)
+    parts = [spec.compute_partial(Bundle.open(rollouts, t, two_ctx, chunk=(k, n_chunks))) for k in range(n_chunks)]
+    chunked = spec.finalize(parts, two_ctx)

    assert chunked.id == unchunked.id
    assert chunked.kind == unchunked.kind
    _assert_payload_close(unchunked.payload, chunked.payload)
+
+
+# ---------------------------------------------------------------------------
+# new (gitea #76) reductions: KS distance, confusion matrix, containment depth
+# ---------------------------------------------------------------------------
+
+
+def test_ks_statistic():
+    assert _ks_statistic([10, 10], [10, 10]) == 0.0  # identical shape -> 0
+    assert _ks_statistic([10, 0], [0, 10]) == 1.0  # fully disjoint -> 1
+    assert _ks_statistic([0, 0], [0, 0]) != _ks_statistic([0, 0], [0, 0])  # nan (no data either side)
+    assert _ks_statistic([10, 0], [0, 0]) == 1.0  # one side empty, other isn't -> maximal mismatch
+
+
+def test_integer_confusion_matches_event_pairing():
+    # true (reference) n_sec = [1, 1]; predicted (rollout) n_sec = [1, 0]
+    labels, mat = _integer_confusion(np.array([1, 1]), np.array([1, 0]))
+    assert labels == ["0", "1+"]
+    assert mat.tolist() == [[0, 0], [1, 1]]  # row=true, col=pred
+
+
+def test_integer_confusion_caps_pathological_outliers():
+    labels, mat = _integer_confusion(np.array([0, 500]), np.array([0, 0]), max_bins=5)
+    assert labels[-1] == "4+"
+    assert mat.shape == (5, 5)
+    assert mat.sum() == 2
+
+
+def test_integer_confusion_explicit_cap_overrides_local_range():
+    # Even though this pair's own max is 1, an explicit shared cap forces a
+    # wider (and so cross-rollout-consistent) label set.
+    labels, mat = _integer_confusion(np.array([1, 1]), np.array([0, 1]), cap=3)
+    assert labels == ["0", "1", "2", "3+"]
+    assert mat.shape == (4, 4)
+
+
+def test_containment_depths_simple_ramp():
+    # one event, edep concentrated in the first bin -> 90%/95% containment
+    # depth is the first bin's right edge; a zero-energy event is dropped.
+    mat = np.array([[9.0, 1.0, 0.0], [0.0, 0.0, 0.0]])
+    edges = np.array([0.0, 1.0, 2.0, 3.0])
+    depths = _containment_depths(mat, edges, 0.90)
+    assert depths.tolist() == [1.0]
+
+
+def test_n_sec_confusion_spec(bundle):
+    spec = get_spec("n_sec_confusion")
+    r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
+    assert r.payload["row_labels"] == r.payload["col_labels"] == ["0", "1+"]
+    assert r.payload["series"]["rollout"] == [[0, 0], [1, 1]]
+
+
+def test_n_sec_confusion_shares_one_cap_across_rollouts(two_bundle):
+    spec = get_spec("n_sec_confusion")
+    r = spec.finalize([spec.compute_partial(two_bundle)], two_bundle.ctx)
+    assert list(r.payload["series"]) == ["flow", "wgan"]
+    # both rollouts share the same fixture data here, so their matrices (and
+    # the shared label set) must be identical.
+    assert r.payload["series"]["flow"] == r.payload["series"]["wgan"]
@@ -0,0 +1,280 @@
+"""Tests for giant.checkpoint_io.load_for_inference (issues.md Issue 5) —
+the shared bootstrap `giant predict`/`giant rollout` use to go from a
+checkpoint path to ready-to-run models."""
+
+from __future__ import annotations
+
+import copy
+
+import numpy as np
+import pytest
+import torch
+
+from giant import config as gconfig
+from giant.checkpoint_io import (
+    CheckpointCompatibilityError,
+    InferenceContext,
+    conditioning_axes,
+    load_for_inference,
+    stage_cfg,
+)
+from giant.data.loader import TopNMap
+from giant.data.setup_cache import topnmap_to_json
+from giant.data.transforms import Normalizer
+from giant.model.network import build_models
+
+PDG_MAP = {11: 0, 22: 1, -11: 2}
+MAT_MAP = {"G4_PbWO4": 0, "G4_AIR": 1}
+
+
+def _model_cfg(stage2_active: bool = True) -> dict:
+    """DEFAULT_CONFIG-derived, shrunk for speed — same pattern as
+    tests/test_network.py::_minimal_model_config. Default `conditioning`
+    (both axes "physical") needs no top-N vocab map, so this is a cheap,
+    fully self-contained happy-path config."""
+    cfg = copy.deepcopy(gconfig.DEFAULT_CONFIG)
+    cfg["conditioning"]["particle"]["emb_dim"] = 4
+    cfg["conditioning"]["material"]["emb_dim"] = 4
+    cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1})
+    cfg["stage2_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "k_max": 3})
+    cfg["stage2_model"]["active"] = stage2_active
+    return {
+        "pdg_vocab": len(PDG_MAP),
+        "mat_vocab": len(MAT_MAP),
+        "conditioning": cfg["conditioning"],
+        "stage1_model": cfg["stage1_model"],
+        "stage2_model": cfg["stage2_model"],
+    }
+
+
+def _norms() -> tuple[Normalizer, Normalizer, Normalizer]:
+    rng = np.random.default_rng(0)
+    cond = Normalizer().fit(rng.standard_normal((100, 15)).astype(np.float32))
+    tgt = Normalizer().fit(rng.standard_normal((100, 9)).astype(np.float32))
+    sec_phys = Normalizer().fit(rng.standard_normal((100, 2)).astype(np.float32))
+    return cond, tgt, sec_phys
+
+
+def _write_checkpoint(tmp_path, model_cfg=None, ema: bool = False, **ckpt_overrides):
+    cfg = model_cfg if model_cfg is not None else _model_cfg()
+    built = build_models(cfg)
+    stage1, stage2 = built["stage1"], built["stage2"]
+    cond, tgt, sec_phys = _norms()
+    ckpt: dict = {
+        "model_config": cfg,
+        "model": stage1.state_dict() if stage1 is not None else {},
+        "sec_decoder": stage2.state_dict() if stage2 is not None else {},
+        "pdg_map": PDG_MAP,
+        "mat_map": MAT_MAP,
+        "normalizer": {"cond": cond.to_dict(), "target": tgt.to_dict(), "sec_phys": sec_phys.to_dict()},
+        "epoch": 3,
+        "best_val_loss": 0.5,
+    }
+    if ema:
+        ckpt["model_ema"] = stage1.state_dict() if stage1 is not None else {}
+        ckpt["sec_decoder_ema"] = stage2.state_dict() if stage2 is not None else {}
+    # DEFAULT_CONFIG's stage2_model.particle_type.target defaults to
+    # "onehot", and giant train's pipeline (gitea #29) now always writes a
+    # sec_type_topn_map in that case — default one in here too, unless a
+    # test explicitly overrides it, so fixtures represent a real, loadable
+    # checkpoint by default rather than exercising the "missing" guard by
+    # accident.
+    particle_type_target = cfg.get("stage2_model", {}).get("particle_type", {}).get("target", "onehot")
+    if particle_type_target == "onehot" and "sec_type_topn_map" not in ckpt_overrides:
+        default_sec_type_topn = TopNMap(class_map=dict(zip(PDG_MAP, range(len(PDG_MAP)))), other_members={})
+        ckpt["sec_type_topn_map"] = topnmap_to_json(default_sec_type_topn)
+    ckpt.update(ckpt_overrides)
+    path = tmp_path / "ckpt.pt"
+    torch.save(ckpt, path)
+    return path
+
+
+def _onehot_model_cfg() -> dict:
+    cfg = _model_cfg()
+    cfg["conditioning"]["particle"]["type"] = "onehot"
+    return cfg
+
+
+# ---------------------------------------------------------------------------
+# Happy path
+# ---------------------------------------------------------------------------
+
+
+def test_happy_path_returns_populated_context(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+    assert isinstance(ctx, InferenceContext)
+    assert ctx.stage1 is not None and ctx.stage2 is not None
+    assert not ctx.stage1.training
+    assert not ctx.stage2.training
+    assert next(ctx.stage1.parameters()).device == torch.device("cpu")
+    assert ctx.pdg_map == PDG_MAP
+    assert ctx.mat_map == MAT_MAP
+    assert all(isinstance(k, int) for k in ctx.pdg_map)
+    assert all(isinstance(k, str) for k in ctx.mat_map)
+    assert ctx.particle_conditioning == "physical"
+    assert ctx.material_conditioning == "physical"
+    assert ctx.k_max == 3
+    assert ctx.epoch == 3
+    assert ctx.best_val_loss == 0.5
+    assert ctx.model_config["stage1_model"]["hidden_dim"] == 8
+
+
+def test_happy_path_normalizer_values_round_trip(tmp_path):
+    cond, tgt, sec_phys = _norms()
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+    assert ctx.cond_norm.mean is not None and cond.mean is not None
+    assert ctx.tgt_norm.mean is not None and tgt.mean is not None
+    assert ctx.sec_phys_norm.mean is not None and sec_phys.mean is not None
+    np.testing.assert_allclose(ctx.cond_norm.mean, cond.mean)
+    np.testing.assert_allclose(ctx.tgt_norm.mean, tgt.mean)
+    np.testing.assert_allclose(ctx.sec_phys_norm.mean, sec_phys.mean)
+
+
+# ---------------------------------------------------------------------------
+# Guards
+# ---------------------------------------------------------------------------
+
+
+def test_missing_model_config_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["model_config"]
+    torch.save(ckpt, checkpoint)
+
+    with pytest.raises(CheckpointCompatibilityError, match="no model_config"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_missing_sec_decoder_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["sec_decoder"]
+    torch.save(ckpt, checkpoint)
+
+    with pytest.raises(CheckpointCompatibilityError, match="no sec_decoder"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_missing_sec_phys_normalizer_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["normalizer"]["sec_phys"]
+    torch.save(ckpt, checkpoint)
+
+    with pytest.raises(CheckpointCompatibilityError, match="no normalizer.sec_phys"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_onehot_particle_conditioning_without_topn_map_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, model_cfg=_onehot_model_cfg())
+
+    with pytest.raises(CheckpointCompatibilityError, match="pdg_topn_map"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_onehot_particle_conditioning_with_topn_map_succeeds(tmp_path):
+    topn = TopNMap(class_map={11: 0, 22: 1}, other_members={})
+    checkpoint = _write_checkpoint(
+        tmp_path,
+        model_cfg=_onehot_model_cfg(),
+        pdg_topn_map=topnmap_to_json(topn),
+    )
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+    assert ctx.particle_conditioning == "onehot"
+    assert ctx.pdg_topn_map is not None
+    assert ctx.pdg_topn_map.class_map == {11: 0, 22: 1}
+
+
+def test_onehot_particle_type_target_without_sec_type_topn_map_raises(tmp_path):
+    """DEFAULT_CONFIG's stage2_model.particle_type.target="onehot" needs a
+    sec_type_topn_map (gitea #29) — a checkpoint with neither key at all
+    (not even the pre-#29 pdg_topn_map to fall back to) must fail loudly."""
+    checkpoint = _write_checkpoint(tmp_path, sec_type_topn_map=None)
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["sec_type_topn_map"]
+    torch.save(ckpt, checkpoint)
+
+    with pytest.raises(CheckpointCompatibilityError, match="sec_type_topn_map"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_pre_gitea_29_checkpoint_falls_back_to_pdg_topn_map_for_sec_type(tmp_path):
+    """A checkpoint written before gitea #29 has no sec_type_topn_map key at
+    all — conditioning and secondary-type onehot maps were always the same
+    map, saved once under pdg_topn_map. load_for_inference must reproduce
+    that exact pre-#29 behavior for such a checkpoint."""
+    topn = TopNMap(class_map={11: 0, 22: 1, -11: 2}, other_members={})
+    checkpoint = _write_checkpoint(
+        tmp_path,
+        model_cfg=_onehot_model_cfg(),
+        pdg_topn_map=topnmap_to_json(topn),
+        sec_type_topn_map=None,
+    )
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["sec_type_topn_map"]
+    torch.save(ckpt, checkpoint)
+
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+    assert ctx.sec_type_topn_map is not None
+    assert ctx.sec_type_topn_map.class_map == {11: 0, 22: 1, -11: 2}
+
+
+def test_ema_weights_requested_but_missing_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, ema=False)
+
+    with pytest.raises(CheckpointCompatibilityError, match="no EMA weights"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict", weights="ema")
+
+
+def test_ema_weights_requested_and_present_succeeds(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, ema=True)
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict", weights="ema")
+    assert ctx.stage1 is not None and ctx.stage2 is not None
+
+
+@pytest.mark.parametrize("command_name", ["predict", "rollout"])
+def test_inactive_stage_with_require_stage2_raises_with_command_name(tmp_path, command_name):
+    checkpoint = _write_checkpoint(tmp_path, model_cfg=_model_cfg(stage2_active=False))
+
+    with pytest.raises(CheckpointCompatibilityError, match=f"{command_name} needs both"):
+        load_for_inference(checkpoint, torch.device("cpu"), command_name)
+
+
+def test_inactive_stage_with_require_stage2_false_succeeds_with_stage2_none(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, model_cfg=_model_cfg(stage2_active=False))
+
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict", require_stage2=False)
+    assert ctx.stage1 is not None
+    assert ctx.stage2 is None
+
+
+# ---------------------------------------------------------------------------
+# conditioning_axes / stage_cfg
+# ---------------------------------------------------------------------------
+
+
+def test_conditioning_axes_v02_flat_string_applies_to_both_axes():
+    assert conditioning_axes({"conditioning": "embedding"}) == ("embedding", "embedding")
+
+
+def test_conditioning_axes_v03_nested_dict_independent_per_axis():
+    model_cfg = {"conditioning": {"particle": {"type": "onehot"}, "material": {"type": "physical"}}}
+    assert conditioning_axes(model_cfg) == ("onehot", "physical")
+
+
+def test_conditioning_axes_missing_key_uses_default():
+    assert conditioning_axes({}, default="embedding") == ("embedding", "embedding")
+
+
+def test_stage_cfg_new_shape_returns_subdict():
+    model_cfg = {"stage2_model": {"k_max": 7}}
+    assert stage_cfg(model_cfg, "stage2") == {"k_max": 7}
+
+
+def test_stage_cfg_v02_flat_shape_returns_empty_dict():
+    model_cfg = {"hidden_dim": 32, "n_blocks": 4}
+    assert stage_cfg(model_cfg, "stage2") == {}
@@ -91,6 +91,31 @@ def test_dry_run_writes_nothing(tmp_path: Path):
    assert not out_dir.exists()


+def test_stage1_init_from_and_freeze_flags_scaffold_a_partial_retrain_config(tmp_path: Path):
+    """gitea #42."""
+    out_dir = tmp_path / "run5"
+    result = runner.invoke(
+        app,
+        [
+            "new-run",
+            "--out",
+            str(out_dir),
+            "--stage1-init-from",
+            "ckpt/stage1_good/best.pt",
+            "--stage1-freeze",
+        ],
+    )
+    assert result.exit_code == 0, result.output
+
+    with open(out_dir / "config.toml", "rb") as f:
+        cfg = tomllib.load(f)
+
+    assert cfg["stage1_model"]["init_from"] == "ckpt/stage1_good/best.pt"
+    assert cfg["stage1_model"]["freeze"] is True
+    assert cfg["stage2_model"]["init_from"] == ""
+    assert cfg["stage2_model"]["freeze"] is False
+
+
 def test_force_guard_refuses_to_clobber_existing_checkpoints(tmp_path: Path):
    out_dir = tmp_path / "run5"
    out_dir.mkdir()
@@ -101,9 +126,7 @@ def test_force_guard_refuses_to_clobber_existing_checkpoints(tmp_path: Path):
    assert "already has last.pt" in result.output
    assert not (out_dir / "config.toml").exists()

-    result = runner.invoke(
-        app, ["new-run", "--out", str(out_dir), "--mode", "ddpm", "--force"]
-    )
+    result = runner.invoke(app, ["new-run", "--out", str(out_dir), "--mode", "ddpm", "--force"])
    assert result.exit_code == 0, result.output
    assert (out_dir / "config.toml").exists()

@@ -1,13 +1,18 @@
 import uuid

+import torch
 import yaml
+from typer.testing import CliRunner

 from giant.cli import (
    _CEPH_PREDICTIONS,
    _resolve_prediction_output,
    _write_prediction_ref,
+    app,
 )

+runner = CliRunner()
+

 # ---------------------------------------------------------------------------
 # _resolve_prediction_output
@@ -116,9 +121,7 @@ def test_ref_yaml_includes_comment_when_provided(tmp_path):
    dataset = tmp_path / "full.manifest"
    pred_uuid = str(uuid.uuid4())

-    ref_path = _write_prediction_ref(
-        checkpoint, pred_uuid, out, dataset, comment="baseline sweep run 3"
-    )
+    ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset, comment="baseline sweep run 3")
    data = yaml.safe_load(ref_path.read_text())

    assert data["comment"] == "baseline sweep run 3"
@@ -133,9 +136,7 @@ def test_ref_timestamp_is_iso_format(tmp_path):
    checkpoint.touch()

    pred_uuid = str(uuid.uuid4())
-    ref_path = _write_prediction_ref(
-        checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d"
-    )
+    ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
    data = yaml.safe_load(ref_path.read_text())

    # Must parse without error and be timezone-aware (UTC).
@@ -150,9 +151,24 @@ def test_ref_checkpoint_path_is_absolute(tmp_path):
    checkpoint.touch()

    pred_uuid = str(uuid.uuid4())
-    ref_path = _write_prediction_ref(
-        checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d"
-    )
+    ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
    data = yaml.safe_load(ref_path.read_text())

    assert data["checkpoint"].startswith("/")
+
+
+# ---------------------------------------------------------------------------
+# Bootstrap failure surfaces via the CLI (issues.md Issue 5 — confirms
+# CheckpointCompatibilityError -> typer.Exit(1) actually wires up end-to-end,
+# not just at the giant.checkpoint_io unit level).
+# ---------------------------------------------------------------------------
+
+
+def test_predict_exits_1_on_checkpoint_missing_model_config(tmp_path):
+    checkpoint = tmp_path / "bad.pt"
+    torch.save({"sec_decoder": {}, "normalizer": {"sec_phys": {}}}, checkpoint)
+
+    result = runner.invoke(app, ["predict", "dummy.parquet", "--checkpoint", str(checkpoint)])
+
+    assert result.exit_code == 1
+    assert "checkpoint has no model_config" in result.output
@@ -0,0 +1,33 @@
+"""Thin CLI smoke coverage for `giant rollout` (issues.md Issue 5) — confirms
+the CheckpointCompatibilityError raised by giant.checkpoint_io.load_for_inference
+surfaces as a clean typer.Exit(1) with the expected message, end-to-end
+through the CLI, not just at the giant.checkpoint_io unit level."""
+
+from __future__ import annotations
+
+import torch
+from typer.testing import CliRunner
+
+from giant.cli import app
+
+runner = CliRunner()
+
+
+def test_rollout_exits_1_on_checkpoint_missing_model_config(tmp_path):
+    checkpoint = tmp_path / "bad.pt"
+    torch.save({"sec_decoder": {}, "normalizer": {"sec_phys": {}}}, checkpoint)
+
+    result = runner.invoke(
+        app,
+        [
+            "rollout",
+            "dummy.parquet",
+            "--checkpoint",
+            str(checkpoint),
+            "--geometry",
+            "dummy_geometry.pkl",
+        ],
+    )
+
+    assert result.exit_code == 1
+    assert "checkpoint has no model_config" in result.output
@@ -1,6 +1,5 @@
-"""Tests for `giant train`'s stage-prefixed CLI flags (docs/v0.3.0-design.md
-decision 7 / docs/v0.3.0-followups.md item 2): --stage1-*/--stage2-* must
-independently override each stage's config block, and must take precedence
+"""Tests for `giant train`'s stage-prefixed CLI flags: --stage1-*/--stage2-*
+must independently override each stage's config block, and must take precedence
 over the older shared flags (--mode/--hidden-dim/--n-critic/... ) that still
 apply the same value to both stages for backward compatibility."""

@@ -42,6 +41,10 @@ def test_stage_prefixed_generator_overrides_shared_mode(monkeypatch, tmp_path):


 def test_stage2_only_knobs(monkeypatch, tmp_path):
+    # --stage2-stage1-context is exercised separately at the overrides-dict
+    # level (test_overrides_from_flags_stage2_only_knobs in test_config.py):
+    # its only non-default value, "sampled", is rejected by validate_config
+    # (issues.md Issue 1), so it can't appear in a full CLI invocation here.
    cfg = _invoke_and_capture_cfg(
        monkeypatch,
        tmp_path,
@@ -54,15 +57,12 @@ def test_stage2_only_knobs(monkeypatch, tmp_path):
            "32",
            "--stage2-context-dim",
            "16",
-            "--stage2-stage1-context",
-            "sampled",
        ],
    )
    assert cfg["stage2_model"]["decoder"] == "one_shot"
    assert cfg["stage2_model"]["k_max"] == 8
    assert cfg["stage2_model"]["hidden_dim"] == 32
    assert cfg["stage2_model"]["context_dim"] == 16
-    assert cfg["stage2_model"]["stage1_context"] == "sampled"
    # untouched stage1 defaults
    assert cfg["stage1_model"]["hidden_dim"] == 256

@@ -95,3 +95,110 @@ def test_wgan_knobs_split_per_stage(monkeypatch, tmp_path):
    assert cfg["stage1_model"]["wgan"]["gp_weight"] == 10.0
    assert cfg["stage2_model"]["wgan"]["n_critic"] == 5
    assert cfg["stage2_model"]["wgan"]["gp_weight"] == 2.5
+
+
+def test_stage1_init_from_and_freeze_flags_land_in_cfg_and_dont_touch_stage2(monkeypatch, tmp_path):
+    """gitea #42: --stage{1,2}-init-from/--stage{1,2}-freeze are stage-scoped
+    only. --stage1-freeze alone would fail validate_config (freeze requires
+    init_from or --resume), so both flags are passed together here."""
+    cfg = _invoke_and_capture_cfg(
+        monkeypatch,
+        tmp_path,
+        ["--stage1-init-from", "ckpt/stage1_good/best.pt", "--stage1-freeze"],
+    )
+    assert cfg["stage1_model"]["init_from"] == "ckpt/stage1_good/best.pt"
+    assert cfg["stage1_model"]["freeze"] is True
+    assert cfg["stage2_model"]["init_from"] == ""
+    assert cfg["stage2_model"]["freeze"] is False
+
+
+def test_stage2_init_from_and_freeze_flags_land_in_cfg_and_dont_touch_stage1(monkeypatch, tmp_path):
+    cfg = _invoke_and_capture_cfg(
+        monkeypatch,
+        tmp_path,
+        ["--stage2-init-from", "ckpt/stage2_good/best.pt", "--stage2-freeze"],
+    )
+    assert cfg["stage2_model"]["init_from"] == "ckpt/stage2_good/best.pt"
+    assert cfg["stage2_model"]["freeze"] is True
+    assert cfg["stage1_model"]["init_from"] == ""
+    assert cfg["stage1_model"]["freeze"] is False
+
+
+def test_batch_size_invalid_string_errors(monkeypatch, tmp_path):
+    monkeypatch.setattr(cli, "run_train_job", lambda *a, **kw: None)
+    result = runner.invoke(
+        cli.app,
+        ["train", "dummy.parquet", "--out", str(tmp_path / "run"), "--batch-size", "not-a-number"],
+    )
+    assert result.exit_code == 1
+    assert "--batch-size must be an integer or 'auto'" in result.output
+
+
+def test_out_dir_resolution_prefers_explicit_out_over_resume(monkeypatch, tmp_path):
+    captured: dict = {}
+
+    def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
+        captured["out_dir"] = out_dir
+
+    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+
+    resume_dir = tmp_path / "resumed_run"
+    resume_dir.mkdir()
+    (resume_dir / "last.pt").touch()
+    explicit_out = tmp_path / "explicit_run"
+
+    result = runner.invoke(
+        cli.app,
+        ["train", "dummy.parquet", "--out", str(explicit_out), "--resume", str(resume_dir / "last.pt")],
+    )
+    assert result.exit_code == 0, result.output
+    assert captured["out_dir"] == explicit_out
+
+
+def test_out_dir_resolution_falls_back_to_resume_parent(monkeypatch, tmp_path):
+    captured: dict = {}
+
+    def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
+        captured["out_dir"] = out_dir
+
+    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+
+    resume_dir = tmp_path / "resumed_run"
+    resume_dir.mkdir()
+    (resume_dir / "last.pt").touch()
+
+    result = runner.invoke(cli.app, ["train", "dummy.parquet", "--resume", str(resume_dir / "last.pt")])
+    assert result.exit_code == 0, result.output
+    assert captured["out_dir"] == resume_dir
+
+
+def test_out_dir_resolution_defaults_when_neither_out_nor_resume_given(monkeypatch, tmp_path):
+    captured: dict = {}
+
+    def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
+        captured["out_dir"] = out_dir
+
+    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.chdir(tmp_path)
+
+    result = runner.invoke(cli.app, ["train", "dummy.parquet"])
+    assert result.exit_code == 0, result.output
+    assert captured["out_dir"] == Path("checkpoints") / cli.gconfig.default_out_dir_name(cli.gconfig.DEFAULT_CONFIG)
+
+
+def test_batch_size_auto_estimates_and_echoes(monkeypatch, tmp_path):
+    captured: dict = {}
+
+    def _fake_run_train_job(*, data, cfg, out_dir, num_workers, **kwargs):
+        captured["batch_size"] = cfg["train"]["batch_size"]
+
+    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr(cli.gconfig, "estimate_batch_size", lambda hidden_dim, n_blocks, device: 123)
+
+    result = runner.invoke(
+        cli.app,
+        ["train", "dummy.parquet", "--out", str(tmp_path / "run"), "--batch-size", "auto"],
+    )
+    assert result.exit_code == 0, result.output
+    assert captured["batch_size"] == 123
+    assert "batch_size: 123 (auto-estimated from free GPU memory)" in result.output
@@ -0,0 +1,87 @@
+import pytest
+from giant.cond_layout import AXIS_TYPES, CondLayout
+from giant.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM
+
+# ── cond_cat column layout ───────────────────────────────────────────────────
+
+
+def test_topn_cols_neither_onehot():
+    layout = CondLayout.from_types("physical", "embedding")
+    assert (layout.particle_topn_col, layout.material_topn_col) == (None, None)
+    assert layout.cat_dim == 2
+
+
+def test_topn_cols_particle_only():
+    layout = CondLayout.from_types("onehot", "physical")
+    assert (layout.particle_topn_col, layout.material_topn_col) == (2, None)
+    assert layout.cat_dim == 3
+
+
+def test_topn_cols_material_only():
+    layout = CondLayout.from_types("physical", "onehot")
+    assert (layout.particle_topn_col, layout.material_topn_col) == (None, 2)
+    assert layout.cat_dim == 3
+
+
+def test_topn_cols_both_onehot_particle_then_material():
+    layout = CondLayout.from_types("onehot", "onehot")
+    assert (layout.particle_topn_col, layout.material_topn_col) == (2, 3)
+    assert layout.cat_dim == 4
+
+
+def test_dense_vocab_cols_are_mode_independent():
+    """Columns 0/1 are always the dense pdg/material index — giant.model.routers
+    reads them without knowing the conditioning mode."""
+    assert (CondLayout.PDG_COL, CondLayout.MAT_COL) == (0, 1)
+    for particle in AXIS_TYPES:
+        for material in AXIS_TYPES:
+            layout = CondLayout.from_types(particle, material)
+            assert layout.particle_topn_col not in (layout.PDG_COL, layout.MAT_COL)
+            assert layout.material_topn_col not in (layout.PDG_COL, layout.MAT_COL)
+
+
+# ── cond_cont slice layout ───────────────────────────────────────────────────
+
+
+def test_cont_slices_tile_cond_cont_exactly():
+    """base / particle_phys / material_phys must partition cond_cont with no
+    gap and no overlap — a gap or overlap is exactly the silent
+    mis-indexing this object exists to prevent."""
+    layout = CondLayout.from_types("physical", "physical")
+    covered = list(range(*layout.base.indices(COND_DIM)))
+    covered += list(range(*layout.particle_phys.indices(COND_DIM)))
+    covered += list(range(*layout.material_phys.indices(COND_DIM)))
+    assert covered == list(range(COND_DIM))
+
+
+def test_cont_slice_widths_match_constants():
+    layout = CondLayout.from_types("embedding", "embedding")
+    assert layout.base == slice(0, COND_DIM_BASE)
+    assert layout.particle_phys.stop - layout.particle_phys.start == PARTICLE_PHYS_DIM
+    assert layout.material_phys.stop - layout.material_phys.start == MATERIAL_PHYS_DIM
+    assert layout.cont_dim == COND_DIM
+
+
+def test_cont_slices_are_mode_independent():
+    """cond_cont is COND_DIM wide in every mode — a non-"physical" axis gets
+    its block zero-filled rather than dropped, so the slices never move."""
+    physical = CondLayout.from_types("physical", "physical")
+    for particle in AXIS_TYPES:
+        for material in AXIS_TYPES:
+            layout = CondLayout.from_types(particle, material)
+            assert layout.base == physical.base
+            assert layout.particle_phys == physical.particle_phys
+            assert layout.material_phys == physical.material_phys
+
+
+# ── validation ───────────────────────────────────────────────────────────────
+
+
+def test_unknown_particle_type_raises():
+    with pytest.raises(ValueError, match="unknown conditioning.particle.type 'bogus'"):
+        CondLayout.from_types("bogus", "physical")
+
+
+def test_unknown_material_type_raises():
+    with pytest.raises(ValueError, match="unknown conditioning.material.type 'bogus'"):
+        CondLayout.from_types("physical", "bogus")
@@ -27,6 +27,163 @@ def test_conditioning_enum_has_onehot():
    }


+# ---------------------------------------------------------------------------
+# Config dataclasses (issues.md Issue 1)
+# ---------------------------------------------------------------------------
+
+
+def test_giant_config_to_dict_matches_default_config():
+    """DEFAULT_CONFIG is generated from GiantConfig().to_dict() (not
+    hand-maintained), so the two cannot structurally drift apart — but this
+    pins the *equality* too, catching e.g. a stray in-place mutation of
+    DEFAULT_CONFIG added elsewhere after import."""
+    assert gconfig.GiantConfig().to_dict() == gconfig.DEFAULT_CONFIG
+
+
+@pytest.mark.parametrize(
+    "cls",
+    [
+        gconfig.ConditioningAxisConfig,
+        gconfig.ConditioningConfig,
+        gconfig.FlowConfig,
+        gconfig.DdpmConfig,
+        gconfig.Stage1WganConfig,
+        gconfig.Stage2WganConfig,
+        gconfig.RouterConfig,
+        gconfig.Stage2RouterConfig,
+        gconfig.TrunkConfig,
+        gconfig.NSecConfig,
+        gconfig.ParticleTypeConfig,
+        gconfig.AutoregressiveConfig,
+        gconfig.HeadConfig,
+        gconfig.Stage1HeadsConfig,
+        gconfig.Stage2HeadsConfig,
+        gconfig.Stage1ModelConfig,
+        gconfig.Stage2ModelConfig,
+        gconfig.TrainConfig,
+        gconfig.GiantConfig,
+    ],
+)
+def test_config_dataclass_from_dict_round_trips_through_to_dict(cls):
+    assert cls.from_dict(cls().to_dict()) == cls()
+    assert cls.from_dict(None) == cls()
+
+
+def test_stage2_model_config_defaults_match_documented_v030_intent():
+    """The two keys issues.md Issue 1 found drifted between DEFAULT_CONFIG
+    and build_models/StageSpec.from_config's own .get(key, default)
+    fallbacks — pinned directly against the dataclass that is now their
+    shared single source of truth."""
+    spec = gconfig.Stage2ModelConfig()
+    assert spec.decoder == "autoregressive"
+    assert spec.particle_type.target == "onehot"
+
+
+def test_trunk_config_defaults_to_resmlp_for_both_stages():
+    """gitea #33: a v0.2-migrated / pre-existing config with no `trunk` key
+    at all must reproduce today's behaviour exactly."""
+    assert gconfig.Stage1ModelConfig().trunk.type == "resmlp"
+    assert gconfig.Stage2ModelConfig().trunk.type == "resmlp"
+    assert gconfig.DEFAULT_CONFIG["stage1_model"]["trunk"]["type"] == "resmlp"
+    assert gconfig.DEFAULT_CONFIG["stage2_model"]["trunk"]["type"] == "resmlp"
+
+
+def test_trunk_config_defaults_block_conditioning_to_add_for_both_stages():
+    """gitea #34: a pre-existing config with no `block_conditioning` key
+    must reproduce today's additive-bias behaviour exactly."""
+    assert gconfig.Stage1ModelConfig().trunk.block_conditioning == "add"
+    assert gconfig.Stage2ModelConfig().trunk.block_conditioning == "add"
+    assert gconfig.DEFAULT_CONFIG["stage1_model"]["trunk"]["block_conditioning"] == "add"
+    assert gconfig.DEFAULT_CONFIG["stage2_model"]["trunk"]["block_conditioning"] == "add"
+
+
+def test_init_from_freeze_default_to_unset_for_both_stages():
+    """gitea #42: a pre-existing config with no init_from/freeze key must
+    reproduce today's from-scratch, always-training behaviour exactly."""
+    assert gconfig.Stage1ModelConfig().init_from == ""
+    assert gconfig.Stage1ModelConfig().freeze is False
+    assert gconfig.Stage2ModelConfig().init_from == ""
+    assert gconfig.Stage2ModelConfig().freeze is False
+    assert gconfig.DEFAULT_CONFIG["stage1_model"]["init_from"] == ""
+    assert gconfig.DEFAULT_CONFIG["stage1_model"]["freeze"] is False
+    assert gconfig.DEFAULT_CONFIG["stage2_model"]["init_from"] == ""
+    assert gconfig.DEFAULT_CONFIG["stage2_model"]["freeze"] is False
+
+
+def test_heads_config_defaults_reproduce_pre_gitea_36_hardcoded_shape():
+    """gitea #36: a pre-existing config with no `heads` key must reproduce
+    today's hardcoded `hidden_dim // 2`, one-hidden-layer architecture
+    exactly."""
+    assert gconfig.Stage1ModelConfig().heads.n_sec.hidden_ratio == 0.5
+    assert gconfig.Stage1ModelConfig().heads.n_sec.depth == 2
+    assert gconfig.Stage2ModelConfig().heads.n_sec.hidden_ratio == 0.5
+    assert gconfig.Stage2ModelConfig().heads.n_sec.depth == 2
+    assert gconfig.Stage2ModelConfig().heads.type.hidden_ratio == 0.5
+    assert gconfig.Stage2ModelConfig().heads.type.depth == 2
+    assert gconfig.DEFAULT_CONFIG["stage1_model"]["heads"]["n_sec"] == {"hidden_ratio": 0.5, "depth": 2}
+    assert gconfig.DEFAULT_CONFIG["stage2_model"]["heads"]["n_sec"] == {"hidden_ratio": 0.5, "depth": 2}
+    assert gconfig.DEFAULT_CONFIG["stage2_model"]["heads"]["type"] == {"hidden_ratio": 0.5, "depth": 2}
+
+
+def test_particle_type_config_n_classes_defaults_to_zero_and_round_trips():
+    """gitea #29: n_classes=0 means "inherit conditioning.particle.emb_dim"
+    — the default must stay 0 so an existing config.toml with no
+    stage2_model.particle_type.n_classes key reproduces pre-#29 behavior."""
+    assert gconfig.ParticleTypeConfig().n_classes == 0
+    spec = gconfig.ParticleTypeConfig.from_dict({"n_classes": 32})
+    assert spec.n_classes == 32
+
+
+def test_particle_type_config_class_weighting_defaults_to_none_and_round_trips():
+    """gitea #44: an existing config.toml with no
+    stage2_model.particle_type.class_weighting key must reproduce the
+    pre-#44 unweighted-CE behavior exactly."""
+    assert gconfig.ParticleTypeConfig().class_weighting == "none"
+    spec = gconfig.ParticleTypeConfig.from_dict({"class_weighting": "inverse_freq"})
+    assert spec.class_weighting == "inverse_freq"
+
+
+def test_router_config_extra_round_trips_composed_axis_keys():
+    d = {"enabled": True, "type": "composed", "axis0_type": "energy", "axis0_n_experts": 4}
+    router = gconfig.RouterConfig.from_dict(d)
+    assert router.enabled is True
+    assert router.extra == {"axis0_type": "energy", "axis0_n_experts": 4}
+    assert router.to_dict()["axis0_type"] == "energy"
+
+
+def test_stage2_router_config_tie_to_stage1_not_leaked_into_extra():
+    router = gconfig.Stage2RouterConfig.from_dict({"tie_to_stage1": True})
+    assert router.tie_to_stage1 is True
+    assert "tie_to_stage1" not in router.extra
+
+
+def test_stage1_router_config_has_no_tie_to_stage1_key():
+    """Stage 1's router schema must not gain stage 2's tie_to_stage1 key —
+    that would change every future run's saved config.toml shape."""
+    assert "tie_to_stage1" not in gconfig.RouterConfig().to_dict()
+
+
+def test_n_sec_config_owner_defaults_to_stage2():
+    n_sec = gconfig.NSecConfig()
+    assert n_sec.owner == "stage2"
+
+
+def test_n_sec_config_owner_round_trips():
+    n_sec = gconfig.NSecConfig.from_dict({"mode": "head", "owner": "stage1"})
+    assert n_sec.owner == "stage1"
+    assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1", "stop_sampling": "greedy"}
+
+
+def test_n_sec_config_stop_sampling_defaults_to_greedy():
+    assert gconfig.NSecConfig().stop_sampling == "greedy"
+
+
+def test_n_sec_config_stop_sampling_round_trips():
+    n_sec = gconfig.NSecConfig.from_dict({"mode": "stop_token", "stop_sampling": "sample"})
+    assert n_sec.stop_sampling == "sample"
+    assert n_sec.to_dict()["stop_sampling"] == "sample"
+
+
 # ---------------------------------------------------------------------------
 # _deep_merge
 # ---------------------------------------------------------------------------
@@ -112,9 +269,7 @@ def test_migrate_config_lambda_nsec_and_lambda_s2():


 def test_migrate_config_wgan_knobs_map_to_both_stages():
-    new = gconfig.migrate_config(
-        {"train": {"n_critic": 3, "gp_weight": 5.0, "critic_lr": 1e-4}}
-    )
+    new = gconfig.migrate_config({"train": {"n_critic": 3, "gp_weight": 5.0, "critic_lr": 1e-4}})
    for stage in ("stage1_model", "stage2_model"):
        assert new[stage]["wgan"]["n_critic"] == 3
        assert new[stage]["wgan"]["gp_weight"] == 5.0
@@ -122,9 +277,7 @@ def test_migrate_config_wgan_knobs_map_to_both_stages():


 def test_migrate_config_model_hidden_dim_n_blocks_dropout_map_to_both_stages():
-    new = gconfig.migrate_config(
-        {"model": {"hidden_dim": 128, "n_blocks": 4, "dropout": 0.2}}
-    )
+    new = gconfig.migrate_config({"model": {"hidden_dim": 128, "n_blocks": 4, "dropout": 0.2}})
    for stage in ("stage1_model", "stage2_model"):
        assert new[stage]["hidden_dim"] == 128
        assert new[stage]["n_res_blocks"] == 4
@@ -132,9 +285,7 @@ def test_migrate_config_model_hidden_dim_n_blocks_dropout_map_to_both_stages():


 def test_migrate_config_emb_dim_and_conditioning_map_to_both_axes():
-    new = gconfig.migrate_config(
-        {"model": {"emb_dim": 32, "conditioning": "embedding"}}
-    )
+    new = gconfig.migrate_config({"model": {"emb_dim": 32, "conditioning": "embedding"}})
    for axis in ("particle", "material"):
        assert new["conditioning"][axis]["emb_dim"] == 32
        assert new["conditioning"][axis]["type"] == "embedding"
@@ -181,11 +332,7 @@ def test_migrate_config_router_copied_to_both_stages_with_tie_to_stage1_false():


 def test_migrate_config_router_nonzero_expert_dims_raises():
-    cfg = {
-        "model": {
-            "router": {"enabled": True, "expert_hidden_dim": 128, "expert_n_blocks": 0}
-        }
-    }
+    cfg = {"model": {"router": {"enabled": True, "expert_hidden_dim": 128, "expert_n_blocks": 0}}}
    try:
        gconfig.migrate_config(cfg)
        assert False, "expected ValueError"
@@ -267,9 +414,7 @@ def test_merge_cli_overrides_nested_override_keeps_siblings():
    assert cfg["stage1_model"]["hidden_dim"] == 256  # untouched sibling section


-def test_merge_cli_overrides_file_then_explicit_override_precedence(
-    tmp_path, monkeypatch
-):
+def test_merge_cli_overrides_file_then_explicit_override_precedence(tmp_path, monkeypatch):
    monkeypatch.setattr(gconfig, "git_hash", lambda: "abc123")
    path = tmp_path / "config.toml"
    _write_toml(
@@ -317,9 +462,7 @@ def test_merge_cli_overrides_warns_on_git_hash_mismatch(tmp_path, monkeypatch, c
    assert "current999" in captured.err


-def test_merge_cli_overrides_no_warning_on_matching_git_hash(
-    tmp_path, monkeypatch, capsys
-):
+def test_merge_cli_overrides_no_warning_on_matching_git_hash(tmp_path, monkeypatch, capsys):
    monkeypatch.setattr(gconfig, "git_hash", lambda: "same123")
    path = tmp_path / "config.toml"
    _write_toml(path, git_hash="same123", extra="[train]\nepochs = 5\n")
@@ -328,9 +471,7 @@ def test_merge_cli_overrides_no_warning_on_matching_git_hash(
    assert capsys.readouterr().err == ""


-def test_merge_cli_overrides_no_warning_when_git_hash_unknown(
-    tmp_path, monkeypatch, capsys
-):
+def test_merge_cli_overrides_no_warning_when_git_hash_unknown(tmp_path, monkeypatch, capsys):
    monkeypatch.setattr(gconfig, "git_hash", lambda: "unknown")
    path = tmp_path / "config.toml"
    _write_toml(path, git_hash="abc123", extra="[train]\nepochs = 5\n")
@@ -339,9 +480,7 @@ def test_merge_cli_overrides_no_warning_when_git_hash_unknown(
    assert capsys.readouterr().err == ""


-def test_merge_cli_overrides_no_warning_when_meta_section_absent(
-    tmp_path, monkeypatch, capsys
-):
+def test_merge_cli_overrides_no_warning_when_meta_section_absent(tmp_path, monkeypatch, capsys):
    monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
    path = tmp_path / "config.toml"
    path.write_text("[train]\nepochs = 5\n")
@@ -351,12 +490,8 @@ def test_merge_cli_overrides_no_warning_when_meta_section_absent(


 def test_merge_cli_overrides_real_default_toml_fixture(monkeypatch):
-    monkeypatch.setattr(
-        gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243"
-    )
-    cfg = gconfig.merge_cli_overrides(
-        gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / "default.toml", {}
-    )
+    monkeypatch.setattr(gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243")
+    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / "default.toml", {})
    assert cfg["stage1_model"]["generator"] == "flow"
    assert cfg["stage1_model"]["hidden_dim"] == 256
    assert cfg["stage2_model"]["hidden_dim"] == 256
@@ -431,10 +566,7 @@ def _cfg_with(**dotted_overrides):


 def test_default_out_dir_name_all_defaults_is_just_the_timestamp():
-    assert (
-        gconfig.default_out_dir_name(gconfig.DEFAULT_CONFIG, now=_NOW)
-        == "20260729_1430"
-    )
+    assert gconfig.default_out_dir_name(gconfig.DEFAULT_CONFIG, now=_NOW) == "20260729_1430"


 def test_default_out_dir_name_stage1_generator_shown_bare_no_prefix():
@@ -499,9 +631,7 @@ def test_default_out_dir_name_router_gumbel_shown_when_enabled():
            "stage1_model.router.gumbel": True,
        }
    )
-    assert (
-        gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8_s1gum"
-    )
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8_s1gum"


 def test_default_out_dir_name_overflow_caps_and_hashes_remainder():
@@ -524,9 +654,7 @@ def test_default_out_dir_name_overflow_caps_and_hashes_remainder():
    # First 6 by priority: stage1_generator, stage2_generator, stage2_decoder,
    # stage2_history, particle_type_target, stage1_router — stage2_router
    # overflows into the hash suffix.
-    assert name.startswith(
-        "20260729_1430_wgan_s2-flow_dec-one_shot_hist-attention_pt-physical_s1r-energy8_+"
-    )
+    assert name.startswith("20260729_1430_wgan_s2-flow_dec-one_shot_hist-attention_pt-physical_s1r-energy8_+")


 def test_default_out_dir_name_overflow_hash_is_deterministic_and_value_sensitive():
@@ -584,9 +712,45 @@ def test_validate_config_embedding_target_passes_with_embedding_conditioning():
    gconfig.validate_config(cfg)  # must not raise


+def test_validate_config_bad_class_weighting_rejected():
+    cfg = _cfg_with(**{"stage2_model.particle_type.class_weighting": "effective_num"})
+    with pytest.raises(ValueError, match="class_weighting"):
+        gconfig.validate_config(cfg)
+
+
+def test_validate_config_class_weighting_requires_onehot_target():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.particle_type.class_weighting": "inverse_freq",
+            "stage2_model.particle_type.target": "physical",
+        }
+    )
+    with pytest.raises(ValueError, match="onehot"):
+        gconfig.validate_config(cfg)
+
+
+def test_validate_config_class_weighting_incompatible_with_wgan_generator():
+    # stage2_model.generator defaults to "wgan" and particle_type.target
+    # defaults to "onehot", so only class_weighting needs overriding here.
+    cfg = _cfg_with(**{"stage2_model.particle_type.class_weighting": "inverse_freq"})
+    with pytest.raises(ValueError, match="wgan"):
+        gconfig.validate_config(cfg)
+
+
+def test_validate_config_class_weighting_passes_with_onehot_and_flow():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.particle_type.class_weighting": "inverse_freq",
+            "stage2_model.particle_type.target": "onehot",
+            "stage2_model.generator": "flow",
+        }
+    )
+    gconfig.validate_config(cfg)  # must not raise
+
+
 def test_validate_config_mixed_particle_material_conditioning_is_valid():
-    """docs/v0.3.0-design.md §3.1: the particle and material conditioning
-    axes are configured independently and may mix freely — e.g. material
+    """The particle and material conditioning axes are configured
+    independently and may mix freely — e.g. material
    "physical" with particle "embedding" — and the data pipeline
    (giant/data/transforms.py) now implements that end-to-end, so
    validate_config must not reject it."""
@@ -628,18 +792,171 @@ def test_validate_config_tie_to_stage1_requires_stage1_active():
        assert "tie_to_stage1" in str(e)


-def test_validate_config_stop_token_not_implemented():
-    cfg = _cfg_with(**{"stage2_model.n_sec.mode": "stop_token"})
+@pytest.mark.parametrize("stage_name", ["stage1_model", "stage2_model"])
+def test_validate_config_freeze_without_init_from_or_resume_rejected(stage_name):
+    cfg = _cfg_with(**{f"{stage_name}.freeze": True})
    try:
        gconfig.validate_config(cfg)
        assert False, "expected ValueError"
    except ValueError as e:
-        assert "stop_token" in str(e)
+        assert "init_from" in str(e)
+        assert "--resume" in str(e)
+
+
+@pytest.mark.parametrize("stage_name", ["stage1_model", "stage2_model"])
+def test_validate_config_freeze_with_init_from_passes(stage_name):
+    cfg = _cfg_with(**{f"{stage_name}.freeze": True, f"{stage_name}.init_from": "ckpt/best.pt"})
+    gconfig.validate_config(cfg)  # must not raise
+
+
+@pytest.mark.parametrize("stage_name", ["stage1_model", "stage2_model"])
+def test_validate_config_freeze_without_init_from_passes_under_resume(stage_name):
+    cfg = _cfg_with(**{f"{stage_name}.freeze": True})
+    gconfig.validate_config(cfg, resume=True)  # must not raise
+
+
+def test_validate_config_stop_token_accepted_under_autoregressive():
+    """DEFAULT_CONFIG's stage2_model.decoder is already "autoregressive"
+    (see test_stage2_model_config_defaults_match_documented_v030_intent), so
+    mode="stop_token" alone must not raise."""
+    cfg = _cfg_with(**{"stage2_model.n_sec.mode": "stop_token"})
+    gconfig.validate_config(cfg)  # must not raise
+
+
+def test_validate_config_stop_token_rejected_under_one_shot():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.n_sec.mode": "stop_token",
+            "stage2_model.decoder": "one_shot",
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "stop_token" in str(e) and "autoregressive" in str(e)
+
+
+def test_validate_config_stop_token_rejected_for_stage1_owner():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.n_sec.mode": "stop_token",
+            "stage2_model.n_sec.owner": "stage1",
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "stop_token" in str(e) and "owner" in str(e)
+
+
+def test_validate_config_bad_stop_sampling_rejected():
+    cfg = _cfg_with(**{"stage2_model.n_sec.stop_sampling": "bogus"})
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "stop_sampling" in str(e)
+
+
+def test_validate_config_default_precision_is_fp32():
+    assert gconfig.DEFAULT_CONFIG["train"]["precision"] == "fp32"
+
+
+def test_validate_config_bf16_precision_accepted():
+    cfg = _cfg_with(**{"train.precision": "bf16"})
+    gconfig.validate_config(cfg)  # no raise
+
+
+@pytest.mark.parametrize("bad", ["fp16", "bogus", ""])
+def test_validate_config_bad_precision_rejected(bad):
+    cfg = _cfg_with(**{"train.precision": bad})
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "precision" in str(e)
+
+
+def test_validate_config_stage1_context_sampled_accepted_with_both_stages_active():
+    """gitea #41: 'sampled' is now implemented, so DEFAULT_CONFIG's
+    stage1_model/stage2_model.active = true (both) must let it through."""
+    cfg = _cfg_with(**{"stage2_model.stage1_context": "sampled"})
+    gconfig.validate_config(cfg)  # must not raise
+
+
+def test_validate_config_bad_stage1_context_rejected():
+    cfg = _cfg_with(**{"stage2_model.stage1_context": "bogus"})
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "stage1_context" in str(e)
+
+
+def test_validate_config_stage1_context_sampled_requires_stage1_active():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.stage1_context": "sampled",
+            "stage1_model.active": False,
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "sampled" in str(e) and "stage1_model.active" in str(e)
+
+
+def test_validate_config_stage1_context_sampled_requires_stage2_active():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.stage1_context": "sampled",
+            "stage2_model.active": False,
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "sampled" in str(e) and "stage2_model.active" in str(e)
+
+
+@pytest.mark.parametrize("key", ["ctx_p_start", "ctx_p_end"])
+@pytest.mark.parametrize("value", [-0.1, 1.1])
+def test_validate_config_ctx_p_out_of_range_rejected(key, value):
+    cfg = _cfg_with(
+        **{
+            "stage2_model.stage1_context": "sampled",
+            f"stage2_model.{key}": value,
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert key in str(e)
+
+
+def test_validate_config_stage1_context_sampled_always_truth_rejected_as_noop():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.stage1_context": "sampled",
+            "stage2_model.ctx_p_start": 1.0,
+            "stage2_model.ctx_p_end": 1.0,
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "ctx_p_start" in str(e) and "ctx_p_end" in str(e)


 def test_validate_config_n_sec_truth_rejected_for_rollout_capable_checkpoint():
-    """docs/v0.3.0-design.md §9: 'n_sec.mode = "truth" is invalid for a
-    rollout-capable checkpoint' — both stages active means giant rollout
+    """'n_sec.mode = "truth" is invalid for a rollout-capable checkpoint' —
+    both stages active means giant rollout
    could load this checkpoint, but 'truth' has no ground truth to draw
    n_sec from at rollout time."""
    cfg = _cfg_with(
@@ -688,6 +1005,31 @@ def test_validate_config_ar_default_markov_always_passes():
    gconfig.validate_config(cfg)  # must not raise


+def test_validate_config_ar_order_energy_desc_passes():
+    """'energy_desc' is the only implemented order — must not raise."""
+    cfg = _cfg_with(
+        **{
+            "stage2_model.decoder": "autoregressive",
+            "stage2_model.autoregressive.order": "energy_desc",
+        }
+    )
+    gconfig.validate_config(cfg)  # must not raise
+
+
+def test_validate_config_ar_order_invalid_value_rejected():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.decoder": "autoregressive",
+            "stage2_model.autoregressive.order": "energy_asc",
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "order" in str(e)
+
+
 def test_validate_config_ar_history_attention_passes():
    """v0.3.0 step 7 implements history='attention' — must not raise."""
    cfg = _cfg_with(
@@ -741,11 +1083,12 @@ def test_validate_config_ar_teacher_forcing_invalid_value_rejected():


 def test_validate_config_ar_checks_skipped_under_one_shot():
-    """history/teacher_forcing values that would fail under AR are irrelevant
-    (and unchecked) when decoder='one_shot'."""
+    """order/history/teacher_forcing values that would fail under AR are
+    irrelevant (and unchecked) when decoder='one_shot'."""
    cfg = _cfg_with(
        **{
            "stage2_model.decoder": "one_shot",
+            "stage2_model.autoregressive.order": "bogus",
            "stage2_model.autoregressive.history": "attention",
            "stage2_model.autoregressive.teacher_forcing": "scheduled",
        }
@@ -753,20 +1096,298 @@ def test_validate_config_ar_checks_skipped_under_one_shot():
    gconfig.validate_config(cfg)  # must not raise


+# ---------------------------------------------------------------------------
+# validate_config_keys / merge_cli_overrides unknown-key rejection
+# ---------------------------------------------------------------------------
+
+
+def test_validate_config_keys_default_config_passes():
+    gconfig.validate_config_keys(gconfig.DEFAULT_CONFIG)  # must not raise
+
+
+def test_validate_config_keys_rejects_unknown_top_level_key():
+    cfg = _cfg_with(**{"bogus_section.foo": 1})
+    try:
+        gconfig.validate_config_keys(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "bogus_section" in str(e)
+
+
+def test_validate_config_keys_rejects_unknown_nested_key_with_close_match_hint():
+    cfg = _cfg_with(**{"stage1_model.n_res_block": 12})
+    try:
+        gconfig.validate_config_keys(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "stage1_model.n_res_block" in str(e)
+        assert "n_res_blocks" in str(e)
+
+
+def test_validate_config_keys_allows_composed_router_axis_keys():
+    cfg = _cfg_with(
+        **{
+            "stage1_model.router.enabled": True,
+            "stage1_model.router.type": "composed",
+            "stage1_model.router.axis0_type": "energy",
+            "stage1_model.router.axis0_n_experts": 4,
+            "stage1_model.router.axis1_type": "pdg",
+            "stage1_model.router.axis1_emb_dim": 8,
+        }
+    )
+    gconfig.validate_config_keys(cfg)  # must not raise
+
+
+def test_validate_config_keys_allows_centers_init():
+    cfg = _cfg_with(**{"stage1_model.router.centers_init": [-1.0, 0.0, 1.0]})
+    gconfig.validate_config_keys(cfg)  # must not raise
+
+
+def test_validate_config_keys_rejects_unrelated_unknown_router_key():
+    cfg = _cfg_with(**{"stage1_model.router.n_expert": 4})  # typo for n_experts
+    try:
+        gconfig.validate_config_keys(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "stage1_model.router.n_expert" in str(e)
+        assert "n_experts" in str(e)
+
+
+def test_validate_config_keys_skips_meta_section():
+    cfg = _cfg_with()
+    cfg["meta"] = {"config_version": 3, "git_hash": "abc123"}
+    gconfig.validate_config_keys(cfg)  # must not raise
+
+
+def test_validate_config_keys_allows_trunk_type():
+    cfg = _cfg_with(**{"stage1_model.trunk.type": "resmlp", "stage2_model.trunk.type": "resmlp"})
+    gconfig.validate_config_keys(cfg)  # must not raise
+
+
+def test_validate_config_keys_rejects_unknown_trunk_key():
+    cfg = _cfg_with(**{"stage1_model.trunk.type_o": "resmlp"})  # typo for type
+    try:
+        gconfig.validate_config_keys(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "stage1_model.trunk.type_o" in str(e)
+        assert "type" in str(e)
+
+
+def test_validate_config_keys_allows_block_conditioning():
+    cfg = _cfg_with(
+        **{
+            "stage1_model.trunk.block_conditioning": "film",
+            "stage2_model.trunk.block_conditioning": "adaln",
+        }
+    )
+    gconfig.validate_config_keys(cfg)  # must not raise
+
+
+def test_validate_config_keys_rejects_unknown_block_conditioning_key():
+    cfg = _cfg_with(**{"stage1_model.trunk.block_conditioning_o": "film"})  # typo
+    try:
+        gconfig.validate_config_keys(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "stage1_model.trunk.block_conditioning_o" in str(e)
+        assert "block_conditioning" in str(e)
+
+
+def test_merge_cli_overrides_rejects_typo_in_toml_file(tmp_path):
+    path = tmp_path / "config.toml"
+    path.write_text("[meta]\nconfig_version = 3\n\n[stage1_model]\nn_res_block = 12\n")
+    try:
+        gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {})
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "n_res_block" in str(e)
+
+
+def test_merge_cli_overrides_rejects_typo_in_cli_overrides():
+    try:
+        gconfig.merge_cli_overrides(
+            gconfig.DEFAULT_CONFIG,
+            None,
+            {"stage1_model": {"n_res_block": 12}},
+        )
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "n_res_block" in str(e)
+
+
+@pytest.mark.parametrize("fixture_name", ["default.toml", "wgan_h128_b4_physical.toml"])
+def test_merge_cli_overrides_real_config_fixtures_pass_key_validation(fixture_name, monkeypatch):
+    monkeypatch.setattr(gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243")
+    gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / fixture_name, {})  # must not raise
+
+
+# ---------------------------------------------------------------------------
+# overrides_from_flags (issues.md Issue 3): the flag -> config-path table
+# shared by `giant train`/`giant new-run`. Each test below pins one
+# precedence rule directly, without CliRunner — see also
+# tests/test_cli_train_overrides.py for the thin end-to-end smoke coverage.
+# ---------------------------------------------------------------------------
+
+
+def test_overrides_from_flags_empty_values_yield_empty_overrides():
+    assert gconfig.overrides_from_flags({}) == {}
+    assert gconfig.overrides_from_flags({"epochs": None, "hidden_dim": None}) == {}
+
+
+def test_overrides_from_flags_train_block_passthrough():
+    overrides = gconfig.overrides_from_flags({"epochs": 5, "lr": 1e-3, "hidden_dim": None})
+    assert overrides == {"train": {"epochs": 5, "lr": 1e-3}}
+
+
+def test_overrides_from_flags_precision_passthrough():
+    overrides = gconfig.overrides_from_flags({"precision": "bf16"})
+    assert overrides == {"train": {"precision": "bf16"}}
+
+
+@pytest.mark.parametrize(
+    ("shorthand", "explicit", "path_key"),
+    [
+        ("hidden_dim", "stage1_hidden_dim", "hidden_dim"),
+        ("n_blocks", "stage1_n_res_blocks", "n_res_blocks"),
+        ("dropout", "stage1_dropout", "dropout"),
+    ],
+)
+def test_overrides_from_flags_stage1_explicit_overrides_shorthand(shorthand, explicit, path_key):
+    overrides = gconfig.overrides_from_flags({shorthand: 1, explicit: 2})
+    assert overrides["stage1_model"][path_key] == 2
+
+
+@pytest.mark.parametrize(
+    ("shorthand", "path_key"),
+    [("hidden_dim", "hidden_dim"), ("n_blocks", "n_res_blocks"), ("dropout", "dropout")],
+)
+def test_overrides_from_flags_stage1_shorthand_alone(shorthand, path_key):
+    overrides = gconfig.overrides_from_flags({shorthand: 7})
+    assert overrides["stage1_model"][path_key] == 7
+
+
+def test_overrides_from_flags_stage2_only_knobs():
+    overrides = gconfig.overrides_from_flags(
+        {
+            "stage2_hidden_dim": 32,
+            "stage2_n_res_blocks": 4,
+            "stage2_dropout": 0.1,
+            "stage2_decoder": "one_shot",
+            "stage2_k_max": 8,
+            "stage2_context_dim": 16,
+            "stage2_stage1_context": "sampled",
+        }
+    )
+    assert overrides["stage2_model"] == {
+        "hidden_dim": 32,
+        "n_res_blocks": 4,
+        "dropout": 0.1,
+        "decoder": "one_shot",
+        "k_max": 8,
+        "context_dim": 16,
+        "stage1_context": "sampled",
+    }
+    assert "stage1_model" not in overrides
+
+
+def test_overrides_from_flags_mode_fans_to_both_stages():
+    overrides = gconfig.overrides_from_flags({"mode": "wgan"})
+    assert overrides["stage1_model"]["generator"] == "wgan"
+    assert overrides["stage2_model"]["generator"] == "wgan"
+
+
+def test_overrides_from_flags_stage1_generator_overrides_mode_for_stage1_only():
+    overrides = gconfig.overrides_from_flags({"mode": "wgan", "stage1_generator": "flow"})
+    assert overrides["stage1_model"]["generator"] == "flow"
+    assert overrides["stage2_model"]["generator"] == "wgan"
+
+
+def test_overrides_from_flags_stage2_generator_overrides_mode_for_stage2_only():
+    overrides = gconfig.overrides_from_flags({"mode": "wgan", "stage2_generator": "flow"})
+    assert overrides["stage1_model"]["generator"] == "wgan"
+    assert overrides["stage2_model"]["generator"] == "flow"
+
+
+def test_overrides_from_flags_emb_dim_sets_both_conditioning_axes():
+    overrides = gconfig.overrides_from_flags({"emb_dim": 24})
+    assert overrides["conditioning"]["particle"]["emb_dim"] == 24
+    assert overrides["conditioning"]["material"]["emb_dim"] == 24
+
+
+def test_overrides_from_flags_conditioning_sets_both_axes_type():
+    overrides = gconfig.overrides_from_flags({"conditioning": "onehot"})
+    assert overrides["conditioning"]["particle"]["type"] == "onehot"
+    assert overrides["conditioning"]["material"]["type"] == "onehot"
+
+
+def test_overrides_from_flags_router_config_only_touches_stage1():
+    overrides = gconfig.overrides_from_flags({"router_config": {"enabled": True, "type": "energy"}})
+    assert overrides["stage1_model"]["router"] == {"enabled": True, "type": "energy"}
+    assert "stage2_model" not in overrides
+
+
+@pytest.mark.parametrize(
+    ("shared", "stage1_specific", "stage2_specific", "path_key"),
+    [
+        ("n_critic", "stage1_n_critic", "stage2_n_critic", "n_critic"),
+        ("gp_weight", "stage1_gp_weight", "stage2_gp_weight", "gp_weight"),
+        ("noise_dim", "stage1_noise_dim", "stage2_noise_dim", "noise_dim"),
+        ("critic_lr", "stage1_critic_lr", "stage2_critic_lr", "critic_lr"),
+    ],
+)
+def test_overrides_from_flags_wgan_knobs_split_per_stage(shared, stage1_specific, stage2_specific, path_key):
+    overrides = gconfig.overrides_from_flags({shared: 5.0, stage1_specific: 3.0})
+    assert overrides["stage1_model"]["wgan"][path_key] == 3.0
+    assert overrides["stage2_model"]["wgan"][path_key] == 5.0
+
+    overrides = gconfig.overrides_from_flags({shared: 5.0, stage2_specific: 2.5})
+    assert overrides["stage1_model"]["wgan"][path_key] == 5.0
+    assert overrides["stage2_model"]["wgan"][path_key] == 2.5
+
+
+@pytest.mark.parametrize(
+    ("stage_flag", "stage_model", "path_key"),
+    [
+        ("stage1_critic_hidden_dim", "stage1_model", "critic_hidden_dim"),
+        ("stage1_critic_n_res_blocks", "stage1_model", "critic_n_res_blocks"),
+        ("stage2_critic_hidden_dim", "stage2_model", "critic_hidden_dim"),
+        ("stage2_critic_n_res_blocks", "stage2_model", "critic_n_res_blocks"),
+    ],
+)
+def test_overrides_from_flags_critic_sizing_is_stage_scoped_only(stage_flag, stage_model, path_key):
+    """critic_hidden_dim/critic_n_res_blocks are architectural per-stage
+    knobs (gitea #28) — unlike n_critic/gp_weight/noise_dim/critic_lr above,
+    there is deliberately no shared alias that fans out to both stages."""
+    overrides = gconfig.overrides_from_flags({stage_flag: 32})
+    assert overrides == {stage_model: {"wgan": {path_key: 32}}}
+
+
+@pytest.mark.parametrize(
+    ("init_from_flag", "freeze_flag", "stage_model"),
+    [
+        ("stage1_init_from", "stage1_freeze", "stage1_model"),
+        ("stage2_init_from", "stage2_freeze", "stage2_model"),
+    ],
+)
+def test_overrides_from_flags_init_from_freeze_is_stage_scoped_only(init_from_flag, freeze_flag, stage_model):
+    """gitea #42: no shared alias — a checkpoint has one set of weights per
+    stage, so "freeze both stages from the same file" has no sensible
+    meaning."""
+    overrides = gconfig.overrides_from_flags({init_from_flag: "ckpt/best.pt", freeze_flag: True})
+    assert overrides == {stage_model: {"init_from": "ckpt/best.pt", "freeze": True}}
+
+
 # ---------------------------------------------------------------------------
 # checkpoint config-mismatch warnings (unchanged surface, still exercised)
 # ---------------------------------------------------------------------------


-def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml(
-    tmp_path, monkeypatch, capsys
-):
+def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml(tmp_path, monkeypatch, capsys):
    monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
    ckpt_path = tmp_path / "best.pt"
    ckpt_path.write_bytes(b"")
-    _write_toml(
-        tmp_path / "config.toml", git_hash="old111", extra="[train]\nepochs = 5\n"
-    )
+    _write_toml(tmp_path / "config.toml", git_hash="old111", extra="[train]\nepochs = 5\n")

    gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)

@@ -776,9 +1397,7 @@ def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml(
    assert "current999" in captured.err


-def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent(
-    tmp_path, monkeypatch, capsys
-):
+def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent(tmp_path, monkeypatch, capsys):
    monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
    ckpt_path = tmp_path / "best.pt"
    ckpt_path.write_bytes(b"")
@@ -787,15 +1406,11 @@ def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent(
    assert capsys.readouterr().err == ""


-def test_warn_if_checkpoint_config_mismatch_no_warning_when_hashes_match(
-    tmp_path, monkeypatch, capsys
-):
+def test_warn_if_checkpoint_config_mismatch_no_warning_when_hashes_match(tmp_path, monkeypatch, capsys):
    monkeypatch.setattr(gconfig, "git_hash", lambda: "same123")
    ckpt_path = tmp_path / "best.pt"
    ckpt_path.write_bytes(b"")
-    _write_toml(
-        tmp_path / "config.toml", git_hash="same123", extra="[train]\nepochs = 5\n"
-    )
+    _write_toml(tmp_path / "config.toml", git_hash="same123", extra="[train]\nepochs = 5\n")

    gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)
    assert capsys.readouterr().err == ""
@@ -0,0 +1,137 @@
+"""Consumed-keys audit (issues.md Issue 5).
+
+`validate_config_keys` (`giant/config.py`) only checks that a config key is
+*declared* — present somewhere in `DEFAULT_CONFIG`, which is generated from
+the frozen dataclasses. It says nothing about whether anything actually
+*reads* the value once parsed. Issues 1, 2 and 4 are three keys that slipped
+through exactly that gap: declared, round-tripped, silently ignored. This
+module walks every leaf path in `DEFAULT_CONFIG` and asserts each is either
+genuinely consumed by the model-building/training/rollout code, or explicitly
+recorded in `_KNOWN_UNUSED` with a reason.
+
+"Consumed" is approximated by static analysis rather than true call-graph
+reachability: for each leaf path's field name, does it appear anywhere in a
+fixed whitelist of source files as a real attribute access, a dict-key-shaped
+string constant, or a function/constructor parameter name (the last of these
+because `Router` subclasses receive their config via `**kwargs` filtered by
+signature — see `giant.model.routers.build_router`)? Docstrings are excluded
+from the string-constant scan so prose mentioning a dotted config path in
+passing can't masquerade as a read of it. This whitelist-based approach is
+deliberately narrower than "anywhere in `giant/`": scanning the whole package
+produces false negatives from unrelated identifier collisions (e.g.
+`giant/analysis/router_gating.py`'s `_top1_shares(..., order: list, ...)`
+parameter would otherwise make `stage2_model.autoregressive.order` read as
+"consumed").
+"""
+
+import ast
+from pathlib import Path
+
+from giant.config import DEFAULT_CONFIG
+from giant.config import leaf_paths as _leaf_paths
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+
+# Files that legitimately consume model_config / training config at
+# build/train/rollout time. Not `giant/cli.py` (a CLI flag existing is not
+# consumption — that's precisely how Issue 1 slipped through), not
+# `giant/config.py` itself (declaring/parsing a field is not reading it), and
+# not `giant/model/_legacy.py` (the protected v0.2 migration surface, which
+# intentionally re-derives old flat keys under old names).
+_CONSUMER_ROOTS = ("giant/model", "giant/training")
+_CONSUMER_FILES = (
+    "giant/sample.py",
+    "giant/pipeline.py",
+    "giant/rollout.py",
+    "giant/checkpoint_io.py",
+    "giant/particles.py",
+    "giant/materials.py",
+)
+_EXCLUDED_FILES = ("giant/model/_legacy.py",)
+
+# Leaf DEFAULT_CONFIG paths that are declared but not (yet) read anywhere in
+# the consumer whitelist above. Each entry must name the issue that tracks
+# it. If a key here starts showing up as consumed, the fix landed and this
+# entry is stale — see test_known_unused_allow_list_has_no_stale_entries.
+_KNOWN_UNUSED = {
+    "stage2_model.autoregressive.order": (
+        "gitea #30 — validate_config now checks order is 'energy_desc', but "
+        "nothing in the build/train/rollout consumer whitelist reads the "
+        "value itself since it's still single-valued"
+    ),
+}
+
+# "lambda" is a Python keyword, so the dataclasses expose the dict key
+# "lambda" as the field `lambda_weight` (giant/config.py:49-50).
+_FIELD_NAME_OVERRIDES = {"lambda": "lambda_weight"}
+
+
+def _field_name(leaf_path: str) -> str:
+    name = leaf_path.rsplit(".", 1)[-1]
+    return _FIELD_NAME_OVERRIDES.get(name, name)
+
+
+def _is_docstring_expr(expr: ast.Expr) -> bool:
+    return isinstance(expr.value, ast.Constant) and isinstance(expr.value.value, str)
+
+
+def _collect_names(source: str, filename: str) -> set[str]:
+    tree = ast.parse(source, filename=filename)
+    docstring_ids = set()
+    for node in ast.walk(tree):
+        if isinstance(node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):
+            body = getattr(node, "body", [])
+            if body and isinstance(body[0], ast.Expr) and _is_docstring_expr(body[0]):
+                docstring_ids.add(id(body[0].value))
+
+    names: set[str] = set()
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Attribute):
+            names.add(node.attr)
+        elif isinstance(node, ast.Constant) and isinstance(node.value, str) and id(node) not in docstring_ids:
+            names.add(node.value)
+        elif isinstance(node, ast.arg):
+            names.add(node.arg)
+        elif isinstance(node, ast.keyword) and node.arg is not None:
+            names.add(node.arg)
+    return names
+
+
+def _consumer_files() -> list[Path]:
+    files: set[Path] = {_REPO_ROOT / f for f in _CONSUMER_FILES}
+    for root in _CONSUMER_ROOTS:
+        files |= set((_REPO_ROOT / root).rglob("*.py"))
+    files -= {_REPO_ROOT / f for f in _EXCLUDED_FILES}
+    return sorted(files)
+
+
+def _consumed_names() -> set[str]:
+    names: set[str] = set()
+    for path in _consumer_files():
+        names |= _collect_names(path.read_text(), str(path))
+    return names
+
+
+def test_every_config_key_is_consumed_or_allow_listed():
+    consumed = _consumed_names()
+    unconsumed = {p for p in _leaf_paths(DEFAULT_CONFIG) if _field_name(p) not in consumed}
+    unexplained = unconsumed - _KNOWN_UNUSED.keys()
+    assert not unexplained, (
+        f"config key(s) {sorted(unexplained)} are declared in DEFAULT_CONFIG "
+        "but not read anywhere in the build/train/rollout consumer files "
+        f"({[str(f.relative_to(_REPO_ROOT)) for f in _consumer_files()]}) — "
+        "either wire the key up, or add it to _KNOWN_UNUSED with a reason "
+        "(see issues.md Issue 5)"
+    )
+
+
+def test_known_unused_allow_list_has_no_stale_entries():
+    consumed = _consumed_names()
+    all_paths = set(_leaf_paths(DEFAULT_CONFIG))
+    stale = {p for p in _KNOWN_UNUSED if p not in all_paths or _field_name(p) in consumed}
+    assert not stale, (
+        f"_KNOWN_UNUSED entry/entries {sorted(stale)} no longer belong on the "
+        "allow-list — either the key was removed from DEFAULT_CONFIG, or it "
+        "is now consumed (the underlying issue was fixed). Remove the stale "
+        "entry/entries."
+    )
@@ -4,7 +4,7 @@ from pathlib import Path

 import pytest

-from scripts import create_root_files
+from giant.tools import create_root_files

 parse_detector_spec = create_root_files.parse_detector_spec
 next_shard_index = create_root_files.next_shard_index
@@ -16,9 +16,7 @@ SimJob = create_root_files.SimJob
 PlanError = create_root_files.PlanError


-def _write_fake_executable(
-    path: Path, *, output_count: int = 1, exit_code: int = 0, sleep: float = 0.0
-) -> Path:
+def _write_fake_executable(path: Path, *, output_count: int = 1, exit_code: int = 0, sleep: float = 0.0) -> Path:
    """Stand-in for run_pbwo4/run_sampling: writes *output_count* .root files
    into its own cwd (so callers can verify each job gets an isolated workdir
    and that the workdir ends up holding *only* the .root output, matching
@@ -89,17 +87,13 @@ def test_next_shard_index_continues_past_existing(tmp_path):

 def test_plan_jobs_rejects_missing_gen(tmp_path):
    with pytest.raises(PlanError):
-        plan_jobs(
-            ["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="gen1"
-        )
+        plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="gen1")


 def test_plan_jobs_rejects_malformed_gen(tmp_path):
    (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
    with pytest.raises(PlanError):
-        plan_jobs(
-            ["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="notagen"
-        )
+        plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="notagen")


 def test_plan_jobs_continues_from_existing_shards(tmp_path):
@@ -108,9 +102,7 @@ def test_plan_jobs_continues_from_existing_shards(tmp_path):
    (gen_dir / "pbwo4" / "shard-000.root").touch()
    (gen_dir / "pbwo4" / "shard-001.root").touch()

-    jobs = plan_jobs(
-        ["pbwo4"], num_files=3, dataset_root=tmp_path, kind="steps", gen="gen1"
-    )
+    jobs = plan_jobs(["pbwo4"], num_files=3, dataset_root=tmp_path, kind="steps", gen="gen1")

    assert [j.shard_index for j in jobs] == [2, 3, 4]
    assert all(j.detector == "pbwo4" and j.config is None for j in jobs)
@@ -320,9 +312,7 @@ def test_run_all_caps_concurrency(tmp_path):
    assert {d.name for d in dests} == {f"shard-{i:03d}.root" for i in range(6)}

    intervals = [json.loads(d.read_text()) for d in dests]
-    events = sorted(
-        [(p["start"], 1) for p in intervals] + [(p["end"], -1) for p in intervals]
-    )
+    events = sorted([(p["start"], 1) for p in intervals] + [(p["end"], -1) for p in intervals])
    concurrent = 0
    peak = 0
    for _, delta in events:
@@ -95,7 +95,7 @@ def _dummy_normalizer(width):

 def test_streaming_dataset_offsets_colliding_event_ids_across_files(tmp_path):
    """Two files that each restart event_id from 0 (one Geant4 job per file,
-    see scripts/steps_to_parquet.py) must not have their same-numbered events
+    see giant/tools/steps_to_parquet.py) must not have their same-numbered events
    collapsed together: every row from every file must show up in exactly one
    of train/val, and the number of distinct events must be the sum across
    files, not the union of raw ids."""
@@ -3,15 +3,15 @@ from typer.testing import CliRunner
 from giant import cli as giant_cli
 from giant.config import Conditioning
 from giant.data import setup_cache
-from scripts import dwarf
-from scripts.dwarf import app
+from giant.tools import dwarf
+from giant.tools.dwarf import app
 from test_pipeline import _make_synthetic_steps

 runner = CliRunner()


 def test_conditioning_enum_shared_across_both_clis():
-    """giant.cli and scripts.dwarf must use the one giant.config.Conditioning
+    """giant.cli and giant.tools.dwarf must use the one giant.config.Conditioning
    enum, not independently redefined copies that could silently drift apart
    on valid --conditioning values."""
    assert dwarf.Conditioning is Conditioning
@@ -51,9 +51,7 @@ def test_convert_rejects_output_with_multiple_files(tmp_path):
 def test_convert_rejects_output_with_parallel_jobs(tmp_path):
    root_file = tmp_path / "shard.root"
    root_file.touch()
-    result = runner.invoke(
-        app, ["convert", str(root_file), "--output", "out.parquet", "--jobs", "2"]
-    )
+    result = runner.invoke(app, ["convert", str(root_file), "--output", "out.parquet", "--jobs", "2"])
    assert result.exit_code != 0
    assert "--output cannot be combined with --jobs > 1" in result.output

@@ -126,6 +124,11 @@ def test_warm_cache_router_process_warms_proc_map(tmp_path):
        [
            "warm-cache",
            str(data),
+            # router.type="process" is incompatible with the default
+            # conditioning.particle.type="physical" (validate_config, now
+            # enforced by warm-cache too — see gitea #59).
+            "--particle-conditioning",
+            "embedding",
            "--router",
            "--router-type",
            "process",
@@ -169,3 +172,63 @@ def test_warm_cache_different_val_fraction_is_separate_entry(tmp_path):
    assert loaded is not None
    assert "valfrac=0.1_seed=0_pcond=physical_mcond=physical" in loaded.normalizers
    assert "valfrac=0.3_seed=0_pcond=physical_mcond=physical" in loaded.normalizers
+
+
+def test_warm_cache_config_warms_particle_type_n_classes(tmp_path):
+    """gitea #59: a config setting stage2_model.particle_type.n_classes away
+    from its 0 (= inherit conditioning.particle.emb_dim) default must warm
+    the pdg top-N map under that n_classes, not the emb_dim default, so a
+    later `giant train --config <same file>` run hits it instead of quietly
+    re-scanning every parquet file."""
+    data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)
+    config_path = tmp_path / "config.toml"
+    config_path.write_text("[meta]\nconfig_version = 3\n\n[stage2_model.particle_type]\nn_classes = 32\n")
+
+    runner.invoke(app, ["warm-cache", str(data), "--config", str(config_path)])
+    result = runner.invoke(app, ["warm-cache", str(data), "--config", str(config_path)])
+
+    assert result.exit_code == 0, result.output
+    assert "pdg top-N map: cache hit" in result.output
+    assert "32 classes" in result.output
+
+
+def test_warm_cache_config_rejects_val_fraction_flag(tmp_path):
+    data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)
+    config_path = tmp_path / "config.toml"
+    config_path.write_text("[meta]\nconfig_version = 3\n")
+
+    result = runner.invoke(
+        app,
+        ["warm-cache", str(data), "--config", str(config_path), "--val-fraction", "0.2"],
+    )
+
+    assert result.exit_code != 0
+    assert "--config" in result.output
+    assert "--val-fraction" in result.output
+
+
+def test_warm_cache_config_rejects_router_flags(tmp_path):
+    data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)
+    config_path = tmp_path / "config.toml"
+    config_path.write_text("[meta]\nconfig_version = 3\n")
+
+    result = runner.invoke(
+        app,
+        [
+            "warm-cache",
+            str(data),
+            "--config",
+            str(config_path),
+            "--router",
+            "--router-type",
+            "process",
+            "--n-experts",
+            "3",
+        ],
+    )
+
+    assert result.exit_code != 0
+    assert "--config" in result.output
+    assert "--router/--no-router" in result.output
+    assert "--router-type" in result.output
+    assert "--n-experts" in result.output
@@ -1,11 +1,12 @@
 import torch
+from giant.config import ConditioningAxisConfig
 from giant.constants import COND_DIM
 from giant.model.network import Stage1Model
 from giant.model.schedule import CosineSchedule, flow_matching_loss
 from giant.sample import sample_flow, sample_ddim

-PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
-MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
+PARTICLE_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
+MATERIAL_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)


 def _small_model():
@@ -4,6 +4,7 @@ from pathlib import Path
 from unittest.mock import patch

 import numpy as np
+import pandas as pd
 import pytest

 from giant import geometry as g
@@ -12,6 +13,70 @@ from giant import geometry as g
 pytest.importorskip("sklearn")


+def _steps_frame(n=5, with_post=True):
+    rng = np.random.default_rng(0)
+    data = {
+        "pre_x": rng.uniform(-10, 10, n),
+        "pre_y": rng.uniform(-10, 10, n),
+        "pre_z": rng.uniform(-10, 10, n),
+        "material": ["G4_AIR"] * n,
+        "layer_id": np.arange(n, dtype=np.int64),
+    }
+    if with_post:
+        data["post_x"] = rng.uniform(-10, 10, n)
+        data["post_y"] = rng.uniform(-10, 10, n)
+        data["post_z"] = rng.uniform(-10, 10, n)
+    return pd.DataFrame(data)
+
+
+def test_iter_point_batches_missing_columns_raises(tmp_path):
+    path = tmp_path / "steps.parquet"
+    pd.DataFrame({"pre_x": [0.0]}).to_parquet(path)
+    with pytest.raises(ValueError, match="missing columns"):
+        next(g._iter_point_batches(path))
+
+
+def test_iter_point_batches_without_post_columns_yields_pre_only(tmp_path):
+    path = tmp_path / "steps.parquet"
+    df = _steps_frame(n=5, with_post=False)
+    df.to_parquet(path)
+
+    (pos, mat, lay) = next(g._iter_point_batches(path))
+
+    assert pos.shape == (5, 3)
+    np.testing.assert_allclose(pos, df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32))
+    assert list(mat) == ["G4_AIR"] * 5
+    np.testing.assert_array_equal(lay, np.arange(5))
+
+
+def test_iter_point_batches_with_post_columns_doubles_and_concatenates_points(
+    tmp_path,
+):
+    path = tmp_path / "steps.parquet"
+    df = _steps_frame(n=5, with_post=True)
+    df.to_parquet(path)
+
+    (pos, mat, lay) = next(g._iter_point_batches(path))
+
+    # Every step contributes both its pre_pos and post_pos, sharing the
+    # step's material/layer_id label — so batches double in length.
+    assert pos.shape == (10, 3)
+    np.testing.assert_allclose(pos[:5], df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32))
+    np.testing.assert_allclose(pos[5:], df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32))
+    assert list(mat) == ["G4_AIR"] * 10
+    np.testing.assert_array_equal(lay, np.concatenate([np.arange(5), np.arange(5)]))
+
+
+def test_iter_point_batches_respects_batch_size(tmp_path):
+    path = tmp_path / "steps.parquet"
+    df = _steps_frame(n=10, with_post=False)
+    df.to_parquet(path, row_group_size=10)
+
+    batches = list(g._iter_point_batches(path, batch_size=4))
+
+    assert [len(pos) for pos, _, _ in batches] == [4, 4, 2]
+
+
 def _box_batch(n, rng):
    """A labelled point cloud: inside a 100mm box -> PbWO4/0, else AIR/-1."""
    pos = rng.uniform(-200, 200, (n, 3)).astype(np.float32)
@@ -137,9 +202,7 @@ def test_slab_classes_discovered():

 def test_slab_query_labels_by_depth():
    orc = _build_slab()
-    pos = np.array(
-        [[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0]]
-    )  # layer 0, gap, layer 1
+    pos = np.array([[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0]])  # layer 0, gap, layer 1
    material, layer_id, escaped = orc.query(pos)
    assert list(material) == ["G4_PbWO4", "G4_AIR", "G4_W"]
    assert list(layer_id) == [0, -1, 1]
@@ -168,9 +231,7 @@ def test_slab_save_load_roundtrip(tmp_path):
    orc.save(p)
    loaded = g.GeometryOracle.load(p)

-    pos = np.array(
-        [[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0], [0.0, 0.0, 1e5]]
-    )
+    pos = np.array([[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0], [0.0, 0.0, 1e5]])
    m0, l0, e0 = orc.query(pos)
    m1, l1, e1 = loaded.query(pos)
    assert (m0 == m1).all() and (l0 == l1).all() and (e0 == e1).all()
@@ -0,0 +1,42 @@
+import pytest
+import torch
+
+from giant.model.layers import build_mlp_head
+
+
+def test_build_mlp_head_depth_1_is_bare_linear():
+    head = build_mlp_head(8, 4, hidden=16, depth=1)
+    assert len(head) == 1
+    assert isinstance(head[0], torch.nn.Linear)
+    assert head[0].in_features == 8
+    assert head[0].out_features == 4
+    out = head(torch.randn(3, 8))
+    assert out.shape == (3, 4)
+
+
+def test_build_mlp_head_depth_2_matches_pre_gitea_36_shape():
+    head = build_mlp_head(8, 4, hidden=16, depth=2)
+    assert len(head) == 3
+    assert isinstance(head[0], torch.nn.Linear)
+    assert head[0].in_features == 8
+    assert head[0].out_features == 16
+    assert isinstance(head[1], torch.nn.SiLU)
+    assert isinstance(head[2], torch.nn.Linear)
+    assert head[2].in_features == 16
+    assert head[2].out_features == 4
+    out = head(torch.randn(5, 8))
+    assert out.shape == (5, 4)
+
+
+def test_build_mlp_head_depth_3_has_extra_hidden_layer():
+    head = build_mlp_head(8, 4, hidden=16, depth=3)
+    assert len(head) == 5
+    widths = [(m.in_features, m.out_features) for m in head if isinstance(m, torch.nn.Linear)]
+    assert widths == [(8, 16), (16, 16), (16, 4)]
+    out = head(torch.randn(2, 8))
+    assert out.shape == (2, 4)
+
+
+def test_build_mlp_head_depth_0_raises():
+    with pytest.raises(ValueError, match="depth"):
+        build_mlp_head(8, 4, hidden=16, depth=0)
--- a/Show More
+++ b/Show More