Merge remote-tracking branch 'origin/master' into fix/issue-83

# Conflicts: # .bumpversion.toml # giant/cli.py # pyproject.toml # uv.lock
chore: update changelog for v0.3.16
2026-08-31 12:28:31 +02:00 · 2026-08-31 10:07:24 +00:00 · 2026-08-31 10:07:23 +00:00 · 2026-08-31 12:01:15 +02:00 · 2026-08-31 11:58:31 +02:00 · 2026-08-31 11:58:03 +02:00
117 changed files with 15305 additions and 4942 deletions
@@ -0,0 +1,17 @@
+[tool.bumpversion]
+current_version = "0.4.0"
+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}"
+pre_commit_hooks = ["uv lock", "git add uv.lock"]
+
+[[tool.bumpversion.files]]
+filename = "pyproject.toml"
+search = "version = \"{current_version}\""
+replace = "version = \"{new_version}\""
@@ -2,10 +2,9 @@ name: CI

 "on":
  push:
-    branches: ["**"]
+    branches: ["master"]
    tags: ["**"]
-  pull_request:
-    branches: [master]
+  pull_request: {}

 env:
  UV_CACHE_DIR: /uv-cache
@@ -88,6 +87,92 @@ jobs:
          name: coverage-report
          path: coverage.xml

+  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"
+            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
+
  sync-version-on-tag:
    name: Sync project version with tag
    if: startsWith(github.ref, 'refs/tags/')
@@ -107,7 +192,7 @@ jobs:
            git config user.name "gitea-actions"
            git config user.email "actions@git.larsbogner.de"
            git add pyproject.toml uv.lock
-            git commit -m "chore: sync project version to tag ${GITHUB_REF_NAME} [skip ci]"
+            git commit -m "chore: sync project version to tag ${GITHUB_REF_NAME}"
            git push origin HEAD:master
            git push origin ":refs/tags/${GITHUB_REF_NAME}"
            git tag -f "${GITHUB_REF_NAME}" HEAD
@@ -115,3 +200,33 @@ jobs:
          else
            echo "Tag version matches project version ($CURRENT_VERSION)"
          fi
+
+  publish-package:
+    name: Publish package to Gitea package registry
+    needs: [ruff-check, ruff-format, type-check, test, sync-version-on-tag]
+    if: startsWith(github.ref, 'refs/tags/')
+    runs-on: ubuntu-latest
+    container:
+      image: docker.gitea.com/runner-images:ubuntu-latest
+      volumes:
+        - /srv/act-runner-cache/uv:/uv-cache
+    steps:
+      # Check out by tag name (not the triggering SHA) since sync-version-on-tag
+      # may have force-moved the tag to a version-corrected commit.
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.ref_name }}
+      - uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: false
+      - run: |
+          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
+          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
+      - run: uv build
+      # CI_TOKEN needs write:package scope (in addition to write:repository,
+      # used elsewhere) for this upload to authenticate.
+      - run: |
+          uv publish \
+            --publish-url "https://git.larsbogner.de/api/packages/lars/pypi" \
+            --username gitea-actions \
+            --password "${{ secrets.CI_TOKEN }}"
@@ -0,0 +1,591 @@
+# Changelog
+
+## [0.3.16] - 2026-08-31
+
+### Changed
+
+- Docs: record analysis_341dfb14 baseline rollout benchmark results
+
+- Feat: add eval-cost benchmark — Geant4 reference vs surrogate rollout timing
+
+## [0.3.15] - 2026-08-28
+
+### Changed
+
+- Perf: defer heavy imports in giant/dwarf CLIs until commands run
+
+## [0.3.14] - 2026-08-28
+
+### Changed
+
+- Ci: give automated commits visible checks, scope CI triggers, publish releases
+
+- Ci: fix pull_request trigger not registering
+
+## [0.3.13] - 2026-08-28
+
+### Added
+
+- Add inference-time model_config overrides with a sampling-key allowlist [gitea #87](https://git.larsbogner.de/lars/giant/issues/87)
+
+## [0.3.12] - 2026-08-28
+
+### Added
+
+- Add sampled n_sec under n_sec.mode = 'head' [gitea #86](https://git.larsbogner.de/lars/giant/issues/86)
+
+## [0.3.11] - 2026-08-26
+
+### Changed
+
+- Feat(analysis): per-step secondary multiplicity plots
+
+## [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,88 @@ 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 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`, `sec_count_per_step_by_species` — the latter also drawing the reference as its own panel) 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`).
+**Baseline benchmark (done, 2026-08-26):** `configs/baseline.toml`'s first full rollout-vs-Geant4 validation (`analysis_341dfb14`, checkpoint `20260814_1743_s2-flow_h512_s2h512_bs36864_ep50/best.pt`, epoch 50/50). Confirms the v0.3.0 pivot fixed the species collapse — zero photon secondaries / hallucinated `-14` muon antineutrinos are both gone (γ at 95% of truth, no `-14` in the top species) — and rules out `conditioning.*.type = "physical"` as the cause, since this checkpoint pairs it with `flow`/no-router and still doesn't collapse. Bulk shower observables are close to Geant4 (total deposited energy +1.9%, containment depth-90%/95% both 0.986×), but steps/event now *over*-shoots by 1.32× (the opposite sign from every pre-v0.3.0 checkpoint), no hadronic/nuclear secondaries are produced at all, and event-to-event energy variance is ~16× too narrow. Writeup: `/home/lars/knowledge-base/experiments/giant-baseline-flow-ar-rollout-validation.md`.

-**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.
+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.

-**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.
+**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`.
+
+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`.
@@ -10,6 +10,7 @@ A conditional generative model that replaces the Geant4 step function: given a p
 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

@@ -42,9 +43,9 @@ A **two-stage model**, checkpointed together. Either stage's outcome can be prod

 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).

-**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 (`--conditioning`) — 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.
+**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.

-**MoE routing** (`--router`, either stage): 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.
+**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

@@ -64,12 +65,24 @@ 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, Stage1Model, Stage2OneShot/Stage2Autoregressive, Router/MoE, CriticModel
+│   │   ├── 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
+│   ├── 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
@@ -79,21 +92,34 @@ giant/
 │   │   ├── 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
@@ -116,8 +142,14 @@ giant/
 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`
 ```

+The `dev` extra pulls in `convert`, `analysis`, `geometry`, `wandb` and `workflow` as well.
+
 `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
@@ -137,10 +169,13 @@ Useful flags on `giant train`:
 - `--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`. v0.2 flat-schema configs and checkpoints load fine (auto-migrated).
+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.

@@ -150,11 +185,31 @@ Config-file-only knobs (no CLI flag — use `--config config.toml`): `stage2_mod
 - `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` needs 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 = "^chore: (bump version|update changelog|sync project version)", 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,139 @@
+# 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, now measured against this exact config rather than
+# extrapolated from the pre-v0.3 field (analysis_341dfb14, best.pt @ epoch
+# 50/50, full writeup: knowledge-base/experiments/
+# giant-baseline-flow-ar-rollout-validation.md). Unlike every pre-v0.3
+# checkpoint (which under-produced steps/event by 1.6-5x), this baseline
+# OVER-produces steps/event by 1.32x (1.86e5 vs Geant4 1.41e5) and
+# under-produces secondaries/event by 0.84x (5.97e4 vs 7.14e4) — the sign on
+# steps flipped with the v0.3 autoregressive pivot, so don't assume it still
+# undershoots. Secondary-species hallucination (zero photons, hallucinated
+# `-14` muon antineutrinos) that broke every prior checkpoint is gone; the
+# remaining species gap is a total absence of hadronic/nuclear secondaries
+# (protons, neutrons, ion recoils), not miscalibration of the ones produced.
+# Total deposited energy/event is +1.9% high but its event-to-event spread is
+# ~16x too narrow (31 MeV vs Geant4's 491 MeV). Per-step deposited energy is
+# the worst per-step marginal (KS 0.179 vs 0.004-0.071 for the others).
+
+[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
@@ -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,9 +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
@@ -84,7 +84,7 @@ def _combined_quantiles(r_vals: np.ndarray, t_vals: np.ndarray, lo_q: float, hi_


 def build_context(
-    rollout: str | Path | pl.LazyFrame,
+    rollouts: list[RolloutSpec],
    reference: str | Path | pl.LazyFrame,
    *,
    n_energy_bins: int = 4,
@@ -94,41 +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")
+    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) for name in RANGED_VARS
+        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()

-    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]:
@@ -140,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,
@@ -165,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()},
        },
    )

@@ -0,0 +1,76 @@
+"""Measured Geant4 (miniCaloSim) per-step eval cost — the reference line for
+``eval_cost_per_step`` in ``catalog.py``.
+
+Mirrors the precedent set by ``runtime_estimate.py``'s ``_COST_MODEL``: a
+constant table measured once on a specific machine and pasted in, with the
+methodology and provenance recorded in this docstring rather than derived at
+runtime (there is no live Geant4 install on the machines that run
+``giant analyze``, and re-measuring per invocation would be both slow and
+noisy — see the module docstring precedent).
+
+**Methodology** (``scratchpad/bench_geant4.py``, a one-off, not a `dwarf`
+subcommand): ``run_pbwo4`` (the default homogeneous-PbWO4 miniCaloSim
+executable, see ``~/Programming/minicalosim``) was timed at 3 beam energies
+(1/10/50 GeV) and **4 event counts each**, converting each run's ROOT output
+to Parquet with ``giant.tools.steps_to_parquet.convert_steps_to_parquet``
+immediately after. Event counts were scaled down as energy rose (100/400/
+1000/2000 at 1 GeV, 30/100/200/300 at 10 GeV, 10/25/45/60 at 50 GeV) to keep
+every run's row count under ~8.1M — a naive 50/200 pair at 50 GeV produces
+~27M steps and OOM'd the conversion step on a 14GB laptop. Per-energy linear
+fits (``t = intercept + slope * n``) separate Geant4's one-time init (physics
+tables, geometry construction) from its true marginal per-event cost — the
+slope, not a naive ``t / n_events`` from a single run, is what feeds
+``sim_us_per_step`` below. The per-step denominator is the produced
+``Steps``-tree/Parquet row count, matching the "physical step" unit
+``giant rollout``'s ``timing.n_physical_rows`` uses on the surrogate side.
+Both stages ran single-threaded (default Geant4 threading), pinned to one
+CPU core.
+
+``sim_us_per_step``/``convert_us_per_step``/``sim_ms_per_event`` below are
+the mean across the 3 energies. With 4 event-count points per energy (up
+from an initial 2-point pass, which had ~80% spread and nonsensical negative
+fitted intercepts at 10/50 GeV — an artifact of extrapolating a 2-point
+line), both quantities are now energy-flat as physically expected:
+``sim_us_per_step`` spread ~5%, ``convert_us_per_step`` spread ~13.5%. Treat
+these as reliable to about that precision.
+
+**Caveat — hardware asymmetry**: this reference is single-core CPU. A
+surrogate rollout's ``timing`` block will typically be measured on a batched
+GPU. The resulting ratio in ``eval_cost_per_step`` is a *deployment* speedup
+(what you'd actually see swapping Geant4 for the surrogate in a production
+pipeline), not a same-hardware or per-FLOP comparison — state this whenever
+quoting the number.
+
+**Staleness**: re-run ``scratchpad/bench_geant4.py`` (and update this file)
+if measured on different hardware, after a miniCaloSim/Geant4 version bump,
+or if this reference is more than a year or two stale.
+"""
+
+from __future__ import annotations
+
+GEANT4_REFERENCE: dict = {
+    "sim_us_per_step": 11.2903,
+    "convert_us_per_step": 11.1014,
+    "sim_ms_per_event": 609.6848,
+    "provenance": {
+        "cpu": "AMD Ryzen 7 PRO 4750U with Radeon Graphics",
+        "geant4_version": "11.4.1",
+        "minicalosim_sha": "ea917da",
+        "measured": "2026-08-31",
+        "energies_gev": [1.0, 10.0, 50.0],
+        "spread_pct_sim": 4.96,
+        "spread_pct_convert": 13.52,
+        "threads": 1,
+    },
+}
+
+
+def geant4_per_step_us() -> dict[str, float]:
+    """Sim / convert / total microseconds per physical step, from ``GEANT4_REFERENCE``."""
+    sim = GEANT4_REFERENCE["sim_us_per_step"]
+    convert = GEANT4_REFERENCE["convert_us_per_step"]
+    return {
+        "sim_us_per_step": sim,
+        "convert_us_per_step": convert,
+        "total_us_per_step": sim + convert,
+    }
@@ -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")
@@ -188,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")
@@ -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)
+#   "unavailable"    plot not applicable to this run (e.g. no MoE checkpoint)


@dataclass
@@ -9,20 +9,55 @@ 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

 import numpy as np
 import plotstyle as ps
+from matplotlib.colors import LogNorm
 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 +68,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,8 +85,8 @@ 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 (the `stage1=flow` +
@@ -71,23 +109,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 +138,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 +153,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,7 +195,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"):
@@ -181,11 +238,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")
@@ -195,68 +258,167 @@ 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"))
+    if r.payload.get("log_y"):
+        ax.set_yscale("log")
    ps.style_legend(ax, title="source")
    return fig


 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 = dict(r.payload["series"])
+    row_labels = r.payload["row_labels"]
+    col_labels = r.payload["col_labels"]
+    # A heatmap-shaped plot is one matrix per rollout, so the reference (when the
+    # comparison has one — the distance scorecard doesn't) becomes one more panel
+    # rather than another line.
+    if r.payload.get("reference") is not None:
+        series["reference"] = r.payload["reference"]
+    names = list(series)
+    norm = LogNorm(vmin=1) if r.payload.get("log_color") else None
    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"),
+            norm=norm,
+            vmin=None if norm else r.payload.get("vmin"),
+            vmax=None if norm else 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


@@ -284,13 +446,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:
@@ -349,7 +519,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"},
            },
@@ -375,15 +545,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)
@@ -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
@@ -203,6 +204,7 @@ _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)",
 }


@@ -217,51 +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": []}
+    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)):
@@ -273,41 +327,36 @@ def compute_router_share_by_pdg(
        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):
@@ -317,17 +366,39 @@ def compute_router_share_by_process(
        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 = (
@@ -78,6 +86,7 @@ _PLOT_META_KEYS = (
    "rollout_seed",
    "n_rows",
    "termination_reason_counts",
+    "timing",
    "model_config",
    "training_epoch",
    "best_val_loss",
@@ -111,8 +120,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 +183,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 +210,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,8 +236,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
@@ -178,7 +253,8 @@ def _rows_per_chunk(rollout: str | Path, reference: str | Path, n_chunks: int) -
        )

    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
@@ -186,20 +262,22 @@ def _rows_per_chunk(rollout: str | Path, reference: str | Path, n_chunks: int) -


 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
@@ -208,8 +286,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"):
@@ -217,19 +295,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),
@@ -244,17 +325,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"?,
+    "timing"?}, ...]``, 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,14 +351,17 @@ def compute_reduced(
        raise ValueError(
            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"),
+            timing=r.get("timing"),
+        )
+        for r in rollouts
+    ]
+    bundle = Bundle.open(rollout_specs, reference, ctx, chunk=(chunk_index, effective_n))
    partial = Partial(
        id=spec_id,
        family=spec.family,
@@ -291,16 +377,24 @@ 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"),
+            "timing": ro["plot_meta"].get("timing"),
+        }
+        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"),
    )


@@ -342,136 +436,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,7 +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
@@ -96,6 +96,10 @@ _COST_MODEL: dict[str, tuple[float, float]] = {
    "sec_count_per_species": (0.0, 4.963e-07),
    "sec_energy": (0.0, 4.727e-07),
    "sec_cos_angle": (0.0, 2.749e-06),
+    # Metadata-only (YAML-carried `timing`, no row scan) — same shape as the
+    # router diagnostics' fixed cost, just cheaper since there's no live
+    # torch checkpoint to load.
+    "eval_cost_per_step": (0.0, 0.0),
 }


@@ -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,49 @@ 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
+    timing: 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
+    # Wall-clock cost of this rollout run (giant.cli's rollout command),
+    # from the rollout YAML — eval_cost_per_step only. None on rollout runs
+    # that predate timing instrumentation.
+    timing: dict | None = None
+
+
 def _check_rollout_metadata(path: Path) -> None:
    """Raise if ``path`` carries coord metadata that isn't the rollout tag.

@@ -108,6 +156,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))
@@ -116,9 +180,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))


@@ -168,3 +241,34 @@ def secondaries(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
            pl.col("sec_dz_list").alias("sdz"),
        )
    )
+
+
+def secondaries_by_step(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
+    """One row per produced secondary, tagged with the step that produced it.
+
+    Canonical columns: ``step_key`` (an opaque struct identifying the emitting
+    step) and ``pdg``. ``secondaries`` deliberately drops that link; the
+    per-step multiplicity plots need it, so this is a separate view rather than
+    extra columns every other consumer would pay for.
+
+    - rollout: a secondary's birth row carries ``parent_id`` and a birth
+      position copied verbatim from the parent step's ``post_pos``, so
+      ``(event_id, parent_id, pre_pos)`` identifies the emitting step exactly —
+      no join against the (large) step frame is needed.
+    - reference: secondaries already live on their parent step's row, so the
+      row index *is* the step key. It is only ever used as a group key inside
+      one chunk's own aggregation, so indices repeating across chunks is
+      harmless.
+    """
+    if side is Side.rollout:
+        return lf.filter((pl.col("generation") > 0) & (pl.col("step_no") == 0)).select(
+            pl.struct("event_id", "parent_id", "pre_x", "pre_y", "pre_z").alias("step_key"),
+            "pdg",
+        )
+    return (
+        lf.select("sec_pdg_list")
+        .with_row_index("_row")
+        .explode("sec_pdg_list")
+        .drop_nulls("sec_pdg_list")
+        .select(pl.struct("_row").alias("step_key"), pl.col("sec_pdg_list").cast(pl.Int64).alias("pdg"))
+    )
@@ -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",
        },
    )
@@ -14,11 +14,18 @@ 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)`.
+
+`load_for_inference`'s `config_overrides` (gitea #87) lets a caller change a
+checkpoint's `model_config` at load time, restricted to
+`giant.config.INFERENCE_OVERRIDES` — the allowlist of keys that only affect
+sampling, never module construction/shapes or the preprocessing normalizers/
+vocab maps were fit under.
 """

 from __future__ import annotations

-from dataclasses import dataclass
+import copy
+from dataclasses import dataclass, field
 from pathlib import Path

 import torch
@@ -29,13 +36,45 @@ 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
+from giant.model.network import _migrate_legacy_model_config, build_models


 class CheckpointCompatibilityError(Exception):
    """Checkpoint is missing something `load_for_inference` needs."""


+def apply_config_overrides(model_cfg: dict, overrides: dict[str, object] | None) -> dict:
+    """Deep-merge dotted-path *overrides* into a checkpoint's `model_config`,
+    validated against `giant.config.INFERENCE_OVERRIDES` — the allowlist of
+    keys that only affect sampling, not module construction/shapes or the
+    preprocessing normalizers/vocab maps were fit under (gitea #87).
+
+    Migrates a v0.2 flat `model_config` to the nested v0.3 shape first: a
+    dotted path like "stage1_model.ddpm.n_steps" would otherwise silently
+    write into a dict that `build_models` still reads as flat (it decides
+    v0.2-vs-v0.3 by `"stage1_model" in model_config`), suppressing migration.
+
+    Raises `CheckpointCompatibilityError` — never a bare `ValueError` or a
+    downstream `load_state_dict` size mismatch — for an unknown/disallowed
+    path or a value that fails its allowlisted check.
+    """
+    if not overrides:
+        return model_cfg
+    cfg = model_cfg if "stage1_model" in model_cfg else _migrate_legacy_model_config(model_cfg)
+    cfg = copy.deepcopy(cfg)
+    for path, value in overrides.items():
+        spec = gconfig.INFERENCE_OVERRIDES.get(path)
+        if spec is None:
+            allowed = ", ".join(sorted(gconfig.INFERENCE_OVERRIDES))
+            raise CheckpointCompatibilityError(f"{path!r} is not an inference-safe override — allowed paths: {allowed}")
+        try:
+            spec.check(path, value)
+        except ValueError as exc:
+            raise CheckpointCompatibilityError(str(exc)) from exc
+        gconfig._set_path(cfg, path, value)
+    return cfg
+
+
 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
@@ -88,6 +127,23 @@ def load_mat_topn_map(ckpt: dict) -> TopNMap | None:
    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."""
@@ -101,6 +157,7 @@ class InferenceContext:
    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
@@ -110,6 +167,7 @@ class InferenceContext:
    model_config: dict
    epoch: int | None
    best_val_loss: float | None
+    config_overrides: dict[str, object] = field(default_factory=dict)


 def load_for_inference(
@@ -118,6 +176,7 @@ def load_for_inference(
    command_name: str,
    weights: str = "raw",
    require_stage2: bool = True,
+    config_overrides: dict[str, object] | None = None,
 ) -> InferenceContext:
    """Load *checkpoint* and reconstruct everything `predict`/`rollout` need
    to run it forward, on *device*, in `eval()` mode.
@@ -130,6 +189,13 @@ def load_for_inference(
    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.
+
+    *config_overrides* deep-merges dotted `model_config` paths (e.g.
+    `{"stage2_model.n_sec.sampling": "sample"}`) before anything is
+    derived from `model_config` or built — see `apply_config_overrides` for
+    the allowlist and validation. Every derived `InferenceContext` field
+    (`other_policy`, `stage{1,2}_ddpm_steps`, the built modules, ...)
+    reflects the overridden config.
    """
    ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
    for key in ("model_config", "sec_decoder"):
@@ -141,7 +207,7 @@ def load_for_inference(

    gconfig.warn_if_checkpoint_config_mismatch(checkpoint)

-    model_cfg = ckpt["model_config"]
+    model_cfg = apply_config_overrides(ckpt["model_config"], config_overrides)
    particle_conditioning, material_conditioning = conditioning_axes(model_cfg)
    pdg_topn_map = load_pdg_topn_map(ckpt)
    mat_topn_map = load_mat_topn_map(ckpt)
@@ -153,6 +219,13 @@ def load_for_inference(
        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")
@@ -196,6 +269,7 @@ def load_for_inference(
        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,
@@ -205,4 +279,5 @@ def load_for_inference(
        model_config=model_cfg,
        epoch=ckpt.get("epoch"),
        best_val_loss=ckpt.get("best_val_loss"),
+        config_overrides=dict(config_overrides) if config_overrides else {},
    )
@@ -1,21 +1,20 @@
+from __future__ import annotations
+
 from collections import Counter
 from datetime import datetime, timezone
 from enum import Enum
 import math
 from pathlib import Path
 import re
-from typing import Optional
+import sys
+from typing import TYPE_CHECKING, Optional, cast
 import uuid as uuid_mod

-import numpy as np
-import yaml
-import torch
 import typer
 from typing_extensions import Annotated

-import pyarrow as pa
-import pyarrow.parquet as pq
-from tqdm import tqdm
+if TYPE_CHECKING:
+    import numpy as np

 from giant import config as gconfig
 from giant.constants import (
@@ -25,29 +24,11 @@ from giant.constants import (
    PREDICT_SCHEMA_VERSION_KEY,
    ROLLOUT_COORD_VALUE,
 )
-from giant.data.loader import (
-    event_id_offset,
-    find_parquet_files,
-    iter_file_chunks,
-    iter_cond_chunks,
-)
-from giant.data.transforms import (
-    build_features,
-    build_cond_features,
-    energy_simplex_decode,
-    inv_local_frame_rotation,
-    inv_log_transform,
-    reconstruct_post_pos,
-)
-from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference
-from giant.geometry import GeometryOracle
-from giant.pipeline import run_train_job
-from giant.rollout import (
-    L1DistCollector,
-    decode_secondary_identity,
-    rollout as run_rollout,
-)
-from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
+
+# giant.materials only pulls in numpy (no torch/pandas), and MATERIAL_PROPERTIES
+# is needed at decoration time below (a Typer option default), so it can't be
+# deferred into a command body like the rest of this module's heavy imports.
+from giant.materials import MATERIAL_PROPERTIES

 app = typer.Typer(no_args_is_help=True)

@@ -140,6 +121,23 @@ def _parse_router_axis_flags(specs: list[str]) -> dict[str, object]:
    return out


+def _parse_set_flags(specs: Optional[list[str]]) -> dict[str, object]:
+    """Parse repeated `--set dotted.path=value` flags into a dict, typing
+    each value with `_coerce_scalar` the same way a TOML file's native types
+    would arrive. Validation against the inference-safe allowlist happens
+    downstream in `giant.checkpoint_io.apply_config_overrides` — this only
+    parses syntax.
+    """
+    out: dict[str, object] = {}
+    for spec in specs or []:
+        path, sep, val = spec.partition("=")
+        if not sep:
+            typer.echo(f"error: --set {spec!r} must be 'dotted.path=value'", err=True)
+            raise typer.Exit(1)
+        out[path] = _coerce_scalar(val)
+    return out
+
+
 def _router_cli_overrides(
    router: bool | None,
    router_type: str | None,
@@ -190,8 +188,19 @@ def _write_prediction_ref(
    out: Path,
    dataset_path: Path,
    comment: str | None = None,
+    explicit_out: bool = False,
 ) -> Path:
-    """Write a YAML sidecar in the checkpoint directory and return its path."""
+    """Write the YAML sidecar and return its path.
+
+    With an explicit ``--out`` the sidecar sits next to the output file as
+    ``out.with_suffix(".yaml")`` — a *deterministic* path, which is what lets
+    a workflow task (``giant/workflow/tasks.py``) declare it as a target.
+    Without one, the historic uuid-named file under the checkpoint directory
+    is kept, so ad-hoc runs and the ``/ceph`` predictions convention are
+    unaffected.
+    """
+    import yaml
+
    ref = {
        "prediction_id": pred_uuid,
        "output": str(out),
@@ -201,11 +210,50 @@ def _write_prediction_ref(
    }
    if comment is not None:
        ref["comment"] = comment
-    ref_path = checkpoint.parent / f"{pred_uuid}.yaml"
+    ref_path = out.with_suffix(".yaml") if explicit_out else checkpoint.parent / f"{pred_uuid}.yaml"
    ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False))
    return ref_path


+def _build_rollout_timing(
+    *,
+    setup_s: float,
+    rollout_s: float,
+    write_s: float,
+    n_rows: int,
+    termination_reason_counts: dict[str, int],
+    n_seed_events: int,
+    device: str,
+    torch_threads: int,
+) -> dict:
+    """Assemble ``giant rollout``'s ``timing`` sidecar block.
+
+    ``n_physical_rows`` excludes the synthetic termination rows (escape/
+    unknown-pdg/energy-cutoff/max-steps markers `giant.rollout` emits but
+    Geant4 never does) so ``us_per_step`` is comparable to
+    ``giant.analysis.geant4_reference``'s per-step Geant4 measurement — see
+    ``giant/analysis/catalog.py``'s ``eval_cost_per_step`` spec.
+    """
+    from giant.analysis.sources import SYNTHETIC_TERMINATION_REASONS
+
+    sample_s = rollout_s - write_s
+    n_synthetic_rows = sum(termination_reason_counts.get(reason, 0) for reason in SYNTHETIC_TERMINATION_REASONS)
+    n_physical_rows = n_rows - n_synthetic_rows
+    return {
+        "setup_s": setup_s,
+        "rollout_s": rollout_s,
+        "write_s": write_s,
+        "sample_s": sample_s,
+        "n_rows": n_rows,
+        "n_physical_rows": n_physical_rows,
+        "us_per_step": (sample_s / n_physical_rows * 1e6) if n_physical_rows else None,
+        "write_us_per_step": (write_s / n_physical_rows * 1e6) if n_physical_rows else None,
+        "ms_per_event": (rollout_s / n_seed_events * 1e3) if n_seed_events else None,
+        "device": device,
+        "torch_threads": torch_threads,
+    }
+
+
@app.callback()
 def _main() -> None:
    """GIANT — Geant4 step-function surrogate."""
@@ -228,7 +276,7 @@ class Stage1Context(str, Enum):


 # Conditioning itself lives in giant.config (imported below as gconfig) —
-# shared with scripts/dwarf.py's Typer commands so the two CLIs can't
+# shared with giant/tools/dwarf.py's Typer commands so the two CLIs can't
 # silently drift apart on the option's valid values.
 Conditioning = gconfig.Conditioning

@@ -459,6 +507,64 @@ def train(
        Optional[float],
        typer.Option("--stage2-critic-lr", help="Overrides --critic-lr for stage 2 only"),
    ] = None,
+    stage1_critic_hidden_dim: Annotated[
+        Optional[int],
+        typer.Option(
+            "--stage1-critic-hidden-dim",
+            help="WGAN-GP (--mode wgan only): critic width for stage 1 (default: same as generator's hidden_dim)",
+        ),
+    ] = None,
+    stage1_critic_n_res_blocks: Annotated[
+        Optional[int],
+        typer.Option(
+            "--stage1-critic-n-res-blocks",
+            help="WGAN-GP (--mode wgan only): critic depth for stage 1 (default: same as generator's n_res_blocks)",
+        ),
+    ] = None,
+    stage2_critic_hidden_dim: Annotated[
+        Optional[int],
+        typer.Option(
+            "--stage2-critic-hidden-dim",
+            help="WGAN-GP (--mode wgan only): critic width for stage 2 (default: same as generator's hidden_dim)",
+        ),
+    ] = None,
+    stage2_critic_n_res_blocks: Annotated[
+        Optional[int],
+        typer.Option(
+            "--stage2-critic-n-res-blocks",
+            help="WGAN-GP (--mode wgan only): critic depth for stage 2 (default: same as generator's n_res_blocks)",
+        ),
+    ] = None,
+    stage1_init_from: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--stage1-init-from",
+            help="Checkpoint .pt to load stage 1's weights from before training starts "
+            "(gitea #42) — combine with --stage1-freeze to retrain stage 2 alone "
+            "against a fixed, known-good stage 1",
+        ),
+    ] = None,
+    stage1_freeze: Annotated[
+        Optional[bool],
+        typer.Option(
+            "--stage1-freeze/--no-stage1-freeze",
+            help="Never update stage 1's weights (requires --stage1-init-from, or --resume)",
+        ),
+    ] = None,
+    stage2_init_from: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--stage2-init-from",
+            help="Checkpoint .pt to load stage 2's weights from before training starts (gitea #42)",
+        ),
+    ] = None,
+    stage2_freeze: Annotated[
+        Optional[bool],
+        typer.Option(
+            "--stage2-freeze/--no-stage2-freeze",
+            help="Never update stage 2's weights (requires --stage2-init-from, or --resume)",
+        ),
+    ] = None,
    val_fraction: Annotated[Optional[float], typer.Option("--val-fraction", "-f")] = None,
    seed: Annotated[
        Optional[int],
@@ -553,8 +659,20 @@ def train(
            "steps (default: 50); per-epoch metrics always log in full",
        ),
    ] = None,
+    precision: Annotated[
+        Optional[str],
+        typer.Option(
+            "--precision",
+            help="Training-step autocast precision: 'fp32' (default) or "
+            "'bf16'. No 'fp16' — see giant.training.amp.resolve_autocast",
+        ),
+    ] = None,
 ) -> None:
    """Train the GIANT surrogate model."""
+    import torch
+
+    from giant.pipeline import run_train_job
+
    batch_size_auto = False
    batch_size_value: Optional[int] = None
    if batch_size is not None:
@@ -588,6 +706,7 @@ def train(
        "wandb_project": wandb_project,
        "wandb_run_name": wandb_run_name,
        "wandb_log_every": wandb_log_every,
+        "precision": precision,
        "hidden_dim": hidden_dim,
        "n_blocks": n_blocks,
        "dropout": dropout,
@@ -619,11 +738,19 @@ def train(
        "stage2_gp_weight": stage2_gp_weight,
        "stage2_noise_dim": stage2_noise_dim,
        "stage2_critic_lr": stage2_critic_lr,
+        "stage1_critic_hidden_dim": stage1_critic_hidden_dim,
+        "stage1_critic_n_res_blocks": stage1_critic_n_res_blocks,
+        "stage2_critic_hidden_dim": stage2_critic_hidden_dim,
+        "stage2_critic_n_res_blocks": stage2_critic_n_res_blocks,
+        "stage1_init_from": str(stage1_init_from) if stage1_init_from is not None else None,
+        "stage1_freeze": stage1_freeze,
+        "stage2_init_from": str(stage2_init_from) if stage2_init_from is not None else None,
+        "stage2_freeze": stage2_freeze,
    }
    overrides = gconfig.overrides_from_flags(flag_values)

    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides)
-    gconfig.validate_config(cfg)
+    gconfig.validate_config(cfg, resume=resume is not None)
    t = cfg["train"]

    _device = torch.device(device) if device else gconfig.auto_device()
@@ -658,6 +785,7 @@ def train(

    typer.echo(f"device:  {_device}")
    typer.echo(f"out_dir: {out_dir}")
+    typer.echo(f"precision: {t['precision']}")

    run_train_job(
        data=data,
@@ -703,6 +831,10 @@ def new_run(
    stage2_k_max: Annotated[Optional[int], typer.Option("--stage2-k-max")] = None,
    stage2_context_dim: Annotated[Optional[int], typer.Option("--stage2-context-dim")] = None,
    stage2_stage1_context: Annotated[Optional[Stage1Context], typer.Option("--stage2-stage1-context")] = None,
+    stage1_init_from: Annotated[Optional[Path], typer.Option("--stage1-init-from")] = None,
+    stage1_freeze: Annotated[Optional[bool], typer.Option("--stage1-freeze/--no-stage1-freeze")] = None,
+    stage2_init_from: Annotated[Optional[Path], typer.Option("--stage2-init-from")] = None,
+    stage2_freeze: Annotated[Optional[bool], typer.Option("--stage2-freeze/--no-stage2-freeze")] = None,
    conditioning: Annotated[Optional[Conditioning], typer.Option("--conditioning")] = None,
    router: Annotated[Optional[bool], typer.Option("--router/--no-router")] = None,
    router_type: Annotated[Optional[str], typer.Option("--router-type")] = None,
@@ -763,6 +895,10 @@ def new_run(
        "stage2_k_max": stage2_k_max,
        "stage2_context_dim": stage2_context_dim,
        "stage2_stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
+        "stage1_init_from": str(stage1_init_from) if stage1_init_from is not None else None,
+        "stage1_freeze": stage1_freeze,
+        "stage2_init_from": str(stage2_init_from) if stage2_init_from is not None else None,
+        "stage2_freeze": stage2_freeze,
        "mode": mode.value if mode is not None else None,
        "stage1_generator": stage1_generator.value if stage1_generator is not None else None,
        "stage2_generator": stage2_generator.value if stage2_generator is not None else None,
@@ -822,6 +958,52 @@ def new_run(
    typer.echo(f"  giant train {data_arg} --config {config_path} --out {run_dir}")


+model_app = typer.Typer(
+    no_args_is_help=True,
+    help="Inspect a resolved model architecture without training.",
+)
+app.add_typer(model_app, name="model")
+
+
+@model_app.command("summary")
+def model_summary(
+    config: Annotated[
+        Optional[Path],
+        typer.Option("--config", "-c", help="TOML config file (default: built-in defaults)"),
+    ] = None,
+    pdg_vocab: Annotated[
+        int,
+        typer.Option(
+            "--pdg-vocab",
+            help="Placeholder PDG vocab size for conditioning.particle.type='embedding' "
+            "or a pdg/process router (no dataset attached to derive the real training vocab)",
+        ),
+    ] = 300,
+    mat_vocab: Annotated[
+        int,
+        typer.Option(
+            "--mat-vocab",
+            help="Placeholder material vocab size for conditioning.material.type='embedding' "
+            "or a process router (default: the number of known materials in giant.materials)",
+        ),
+    ] = len(MATERIAL_PROPERTIES),
+) -> None:
+    """Build the resolved model graph from a config with no dataset attached, and
+    print per-module parameter counts, trunk widths, which heads exist, and which
+    conditioning/stage1_model/stage2_model config keys actually shaped it."""
+    from giant.model.summary import render_summary, summarize_model
+
+    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, {})
+    try:
+        gconfig.validate_config(cfg)
+    except ValueError as exc:
+        typer.echo(f"error: {exc}", err=True)
+        raise typer.Exit(1)
+
+    summary = summarize_model(cfg, pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
+    typer.echo(render_summary(summary))
+
+
@app.command()
 def predict(
    data: Annotated[Path, typer.Argument(help="Parquet file or directory of parquet files")],
@@ -889,8 +1071,36 @@ def predict(
            help="Free-text note recorded in the prediction's YAML sidecar",
        ),
    ] = None,
+    set_: Annotated[
+        Optional[list[str]],
+        typer.Option(
+            "--set",
+            help="Override a sampling-only model_config key on this checkpoint, "
+            "'dotted.path=value' (repeatable) — see giant.config.INFERENCE_OVERRIDES "
+            "for the allowlist, e.g. --set stage2_model.n_sec.sampling=sample",
+        ),
+    ] = None,
 ) -> None:
    """Run trained model on a parquet file and save predictions."""
+    import numpy as np
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+    import torch
+    from tqdm import tqdm
+
+    from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference
+    from giant.data.loader import event_id_offset, find_parquet_files, iter_cond_chunks, iter_file_chunks
+    from giant.data.transforms import (
+        build_cond_features,
+        build_features,
+        energy_simplex_decode,
+        inv_local_frame_rotation,
+        inv_log_transform,
+        reconstruct_post_pos,
+    )
+    from giant.rollout import decode_secondary_identity
+    from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
+
    batch_size_auto = False
    batch_size_value: Optional[int] = None
    if batch_size.strip().lower() == "auto":
@@ -909,8 +1119,11 @@ def predict(
    typer.echo(f"device: {_device}")

    # --- Load checkpoint ---
+    config_overrides = _parse_set_flags(set_)
    try:
-        ctx = load_for_inference(checkpoint, _device, "predict", weights=weights.value)
+        ctx = load_for_inference(
+            checkpoint, _device, "predict", weights=weights.value, config_overrides=config_overrides
+        )
    except CheckpointCompatibilityError as exc:
        typer.echo(f"error: {exc}", err=True)
        raise typer.Exit(1)
@@ -921,6 +1134,7 @@ def predict(
    cond_norm, tgt_norm, sec_phys_norm = ctx.cond_norm, ctx.tgt_norm, ctx.sec_phys_norm
    pdg_map, mat_map = ctx.pdg_map, ctx.mat_map
    pdg_topn_map, mat_topn_map = ctx.pdg_topn_map, ctx.mat_topn_map
+    sec_type_topn_map = ctx.sec_type_topn_map
    particle_conditioning, material_conditioning = ctx.particle_conditioning, ctx.material_conditioning
    other_policy = ctx.other_policy
    stage1_ddpm_steps = ctx.stage1_ddpm_steps
@@ -944,6 +1158,7 @@ def predict(
    bs = batch_size_value

    # --- Output path ---
+    explicit_out = out is not None
    out, dataset_path, pred_uuid = _resolve_prediction_output(data, out)
    out.parent.mkdir(parents=True, exist_ok=True)
    typer.echo(f"output: {out}")
@@ -977,7 +1192,7 @@ def predict(
        nonlocal writer, total

        if coord == Coord.local:
-            cond_cont, cond_cat, target_raw, _, _, _, _, _, _ = build_features(
+            feats = build_features(
                piece,
                pdg_map,
                mat_map,
@@ -987,7 +1202,9 @@ def predict(
                mat_topn_map=cond_mat_topn,
                k_max=stage2_k_max,
            )
-            cond_cont = cond_norm.transform(cond_cont)
+            cond_cat = feats.cond_cat
+            target_raw = feats.target_s1
+            cond_cont = cond_norm.transform(feats.cond_cont)
        else:
            cond_cont, cond_cat = build_cond_features(
                piece,
@@ -1008,8 +1225,12 @@ def predict(
            # A fresh v0.3.0 Stage1Model owns no n_sec_head —
            # sample_stage1 returns n_sec_pred=None then, so ask stage 2.
            n_sec_pred = resolve_n_sec(model, sec_decoder, cc, ck, stage1_norm, n_sec_pred)
-            sec_cont, sec_type, _sec_valid_pred = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
-            n_sec_pred_np = n_sec_pred.cpu().numpy()
+            sec_cont, sec_type, sec_valid_pred = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
+            # 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()

        pred = stage1_norm.cpu().numpy()  # normalised

@@ -1071,7 +1292,7 @@ def predict(
                piece["pre_dir"],
                sec_phys_norm,
                pdg_map,
-                pdg_topn_map,
+                sec_type_topn_map,
                other_policy,
                None,
            )
@@ -1157,7 +1378,7 @@ def predict(
    if writer is not None:
        writer.close()

-    ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path, comment)
+    ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path, comment, explicit_out=explicit_out)
    typer.echo(f"reference: {ref_path}")

    if skipped:
@@ -1176,6 +1397,10 @@ def _seed_from_data(files: list[Path], n_events: int | None) -> dict[str, np.nda
    the codebase's convention for the primary (a secondary always carries less
    energy than its parent). See giant/analysis/reduce.py:entry_axis.
    """
+    import numpy as np
+
+    from giant.data.loader import event_id_offset, iter_cond_chunks
+
    best_E: dict[int, float] = {}
    best: dict[int, tuple] = {}
    for file_idx, path in enumerate(files):
@@ -1269,8 +1494,32 @@ def rollout(
        Optional[int],
        typer.Option("--seed", help="Torch/numpy seed for reproducibility"),
    ] = None,
+    set_: Annotated[
+        Optional[list[str]],
+        typer.Option(
+            "--set",
+            help="Override a sampling-only model_config key on this checkpoint, "
+            "'dotted.path=value' (repeatable) — see giant.config.INFERENCE_OVERRIDES "
+            "for the allowlist, e.g. --set stage2_model.n_sec.sampling=sample",
+        ),
+    ] = None,
 ) -> None:
    """Roll the surrogate forward into full showers (autoregressive)."""
+    import time
+
+    import numpy as np
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+    import torch
+    import yaml
+
+    from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference
+    from giant.data.loader import find_parquet_files
+    from giant.geometry import GeometryOracle
+    from giant.rollout import L1DistCollector, RolloutSummary, rollout as run_rollout
+
+    _t_setup_start = time.perf_counter()
+
    if seed is not None:
        torch.manual_seed(seed)
        np.random.seed(seed)
@@ -1278,8 +1527,11 @@ def rollout(
    _device = torch.device(device) if device else gconfig.auto_device()
    typer.echo(f"device: {_device}")

+    config_overrides = _parse_set_flags(set_)
    try:
-        ctx = load_for_inference(checkpoint, _device, "rollout", weights=weights.value)
+        ctx = load_for_inference(
+            checkpoint, _device, "rollout", weights=weights.value, config_overrides=config_overrides
+        )
    except CheckpointCompatibilityError as exc:
        typer.echo(f"error: {exc}", err=True)
        raise typer.Exit(1)
@@ -1292,6 +1544,7 @@ def rollout(
    cond_norm, tgt_norm, sec_phys_norm = ctx.cond_norm, ctx.tgt_norm, ctx.sec_phys_norm
    pdg_map, mat_map = ctx.pdg_map, ctx.mat_map
    pdg_topn_map, mat_topn_map = ctx.pdg_topn_map, ctx.mat_topn_map
+    sec_type_topn_map = ctx.sec_type_topn_map
    particle_conditioning, material_conditioning = ctx.particle_conditioning, ctx.material_conditioning
    other_policy = ctx.other_policy
    stage1_ddpm_steps, stage2_ddpm_steps = ctx.stage1_ddpm_steps, ctx.stage2_ddpm_steps
@@ -1304,6 +1557,7 @@ def rollout(
    seeds = _seed_from_data(files, n_events)
    typer.echo(f"seeded {len(seeds['event_id']):,} shower(s)")

+    explicit_out = out is not None
    out, dataset_path, pred_uuid = _resolve_prediction_output(data, out)
    out.parent.mkdir(parents=True, exist_ok=True)

@@ -1312,9 +1566,11 @@ def rollout(
    # avg_tracks_per_event) — mirrors the row-group streaming `giant predict`
    # already does on its input side.
    writer: pq.ParquetWriter | None = None
+    _write_s = 0.0

    def _write_chunk(row: dict[str, np.ndarray]) -> None:
-        nonlocal writer
+        nonlocal writer, _write_s
+        _t0 = time.perf_counter()
        table = pa.table(row)
        if writer is None:
            table = table.replace_schema_metadata(
@@ -1325,11 +1581,14 @@ def rollout(
            )
            writer = pq.ParquetWriter(out, table.schema)
        writer.write_table(table)
+        _write_s += time.perf_counter() - _t0

    # Only meaningful under particle_type.target="embedding" — a
    # no-op collector otherwise, cheaper than branching the call itself.
    l1_dist_collector = L1DistCollector()

+    _setup_s = time.perf_counter() - _t_setup_start
+    _t_rollout_start = time.perf_counter()
    summary = run_rollout(
        model,
        sec_decoder,
@@ -1352,6 +1611,7 @@ def rollout(
        material_conditioning=material_conditioning,
        pdg_topn_map=pdg_topn_map,
        mat_topn_map=mat_topn_map,
+        sec_type_topn_map=sec_type_topn_map,
        other_policy=other_policy,
        seed=seed,
        stage1_ddpm_steps=stage1_ddpm_steps,
@@ -1360,10 +1620,27 @@ def rollout(
    )
    if writer is not None:
        writer.close()
+    # on_chunk=_write_chunk is always passed above, so rollout() always
+    # returns the streaming-summary shape (RolloutSummary), never the
+    # materialized dict[str, np.ndarray] alternative its return type allows.
+    summary = cast(RolloutSummary, summary)
+    _rollout_s = time.perf_counter() - _t_rollout_start
+    timing = _build_rollout_timing(
+        setup_s=_setup_s,
+        rollout_s=_rollout_s,
+        write_s=_write_s,
+        n_rows=summary["n_rows"],
+        termination_reason_counts=summary["termination_reason_counts"],
+        n_seed_events=len(seeds["event_id"]),
+        device=str(_device),
+        torch_threads=torch.get_num_threads(),
+    )
+    _sample_s = timing["sample_s"]
+    n_physical_rows = timing["n_physical_rows"]

    l1_summary = l1_dist_collector.summary()

-    ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path)
+    ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path, explicit_out=explicit_out)
    ref = yaml.safe_load(ref_path.read_text())
    ref.update(
        {
@@ -1382,6 +1659,10 @@ def rollout(
            "rollout_seed": seed,
            "n_rows": summary["n_rows"],
            "termination_reason_counts": summary["termination_reason_counts"],
+            # Wall-clock cost of this run, normalized per physical step (the
+            # comparable unit against giant.analysis.geant4_reference) — see
+            # eval_cost_per_step in giant/analysis/catalog.py.
+            "timing": timing,
            # Diagnostic — only present under
            # stage2_model.particle_type.target="embedding"; omitted (not
            # written as null) otherwise, so giant.analysis can tell "not
@@ -1392,6 +1673,7 @@ def rollout(
            # model knob (router type/n_experts, noise_dim, vocab sizes, ...)
            # is available downstream without touching this command again.
            "model_config": dict(model_cfg),
+            "config_overrides": dict(ctx.config_overrides),
            "training_epoch": ctx.epoch,
            "best_val_loss": ctx.best_val_loss,
            # [train]/[meta] from the sibling config.toml (giant.config.save_config)
@@ -1404,9 +1686,57 @@ def rollout(

    typer.echo(f"wrote {summary['n_rows']:,} step rows → {out}")
    typer.echo(f"terminations: {summary['termination_reason_counts']}")
+    if timing["us_per_step"] is not None:
+        typer.echo(
+            f"timing: {_rollout_s:.1f}s total ({_sample_s:.1f}s sample + {_write_s:.1f}s write), "
+            f"{timing['us_per_step']:.1f} us/step over {n_physical_rows:,} physical steps"
+        )
    typer.echo(f"reference: {ref_path}")


+workflow_app = typer.Typer(
+    no_args_is_help=True,
+    help="b2luigi pipeline orchestration: one spec file -> cache-warm, train, rollout, analysis.",
+)
+app.add_typer(workflow_app, name="workflow")
+
+
+@workflow_app.command("run")
+def workflow_run(
+    spec: Annotated[Path, typer.Argument(help="Workflow TOML (see configs/workflow_example.toml)")],
+    batch: Annotated[
+        bool,
+        typer.Option("--batch/--local", help="Submit batch-system tasks to HTCondor, or run everything locally"),
+    ] = False,
+    workers: Annotated[int, typer.Option("--workers", help="Concurrent luigi workers")] = 1,
+    mode: Annotated[
+        str,
+        typer.Option(
+            "--mode",
+            help="run | dry-run (print pending tasks) | show-output (print every target) | remove (delete outputs)",
+        ),
+    ] = "run",
+    scheduler_host: Annotated[Optional[str], typer.Option("--scheduler-host", help="luigid host")] = None,
+    scheduler_port: Annotated[Optional[int], typer.Option("--scheduler-port", help="luigid port")] = None,
+) -> None:
+    """Run a workflow spec end to end (the only sanctioned multi-step entry point).
+
+    A thin exec of `giant/workflow/run.py`, which b2luigi also re-executes on
+    every worker — so there is one documented entry point and one code path.
+    """
+    import subprocess
+
+    script = Path(__file__).resolve().parent / "workflow" / "run.py"
+    cmd = [sys.executable, str(script), "--spec", str(spec), "--workers", str(workers), "--mode", mode]
+    if batch:
+        cmd.append("--batch")
+    if scheduler_host:
+        cmd += ["--scheduler-host", scheduler_host]
+    if scheduler_port:
+        cmd += ["--scheduler-port", str(scheduler_port)]
+    raise typer.Exit(subprocess.run(cmd).returncode)
+
+
 analyze_app = typer.Typer(
    no_args_is_help=True,
    help="Rollout-vs-reference analysis: parallel compute on HTCondor + local render.",
@@ -1416,10 +1746,23 @@ app.add_typer(analyze_app, name="analyze")

@analyze_app.command("prep")
 def analyze_prep(
-    rollout_yaml: Annotated[
-        Path,
-        typer.Argument(help="giant rollout YAML sidecar (names the rollout + reference files)"),
+    rollout_yamls: Annotated[
+        list[Path],
+        typer.Argument(
+            help="giant rollout YAML sidecar(s) (names the rollout + reference files). "
+            "Multiple compare N rollouts against one shared reference — every YAML must "
+            "name the same `dataset`."
+        ),
    ],
+    label: Annotated[
+        Optional[list[str]],
+        typer.Option(
+            "--label",
+            help="Series name for a rollout YAML, positionally matched to it — give none, "
+            'or exactly one per YAML. Defaults to the YAML stem (or "rollout" for a '
+            "single YAML).",
+        ),
+    ] = None,
    run_dir: Annotated[
        Optional[Path],
        typer.Option(
@@ -1436,14 +1779,15 @@ def analyze_prep(
        typer.Option("--chunks", help="Split each plot's data into this many event_id chunks"),
    ] = 1,
 ) -> None:
-    """Read the rollout YAML → shared.json + run_meta.json in the run directory."""
+    """Read the rollout YAML(s) → shared.json + run_meta.json in the run directory."""
    from giant.analysis import prep

    path = prep(
-        rollout_yaml,
+        rollout_yamls,
        run_dir,
        n_chunks=chunks,
        default_base=Path.cwd() / "analysis_runs",
+        labels=label,
        n_energy_bins=n_energy_bins,
        n_marginal_bins=n_marginal_bins,
        top_k_pdg=top_k_pdg,
@@ -1504,66 +1848,23 @@ def analyze_render(
    typer.echo(f"rendered {len(pdfs)} plots → {Path(run_dir) / 'plots'}")


-@analyze_app.command("submit")
-def analyze_submit(
-    rollout_yaml: Annotated[Path, typer.Argument(help="giant rollout YAML sidecar")],
-    accounting_group: Annotated[str, typer.Option("--accounting-group")],
-    run_dir: Annotated[
+@analyze_app.command("metrics")
+def analyze_metrics(
+    run_dir: Annotated[Path, typer.Argument(help="Run directory containing metrics.csv (from `giant train`)")],
+    out_dir: Annotated[
        Optional[Path],
        typer.Option(
-            "--run-dir",
+            "--out",
            "-o",
-            help="Override the run directory (default: <cwd>/analysis_runs/analysis_<id>)",
+            help="Override the output directory (default: <cwd>/analysis_runs/metrics_<run_dir name>)",
        ),
    ] = None,
-    docker_image: Annotated[str, typer.Option("--docker-image")] = "cverstege/alma9-gridjob",
-    request_memory: Annotated[int, typer.Option("--request-memory", help="MB")] = 8192,
-    remote: Annotated[
-        bool,
-        typer.Option("--remote/--local", help="+RemoteJob vs ProvidesETPResources"),
-    ] = False,
-    chunks: Annotated[
-        int,
-        typer.Option(
-            "--chunks",
-            help="Split each plot's data into this many event_id chunks/jobs",
-        ),
-    ] = 1,
-    n_energy_bins: Annotated[int, typer.Option("--energy-bins")] = 4,
-    n_marginal_bins: Annotated[int, typer.Option("--bins")] = 50,
-    top_k_pdg: Annotated[int, typer.Option("--top-pdg")] = 6,
-    dry_run: Annotated[bool, typer.Option("--dry-run", help="Write files but don't condor_submit")] = False,
 ) -> None:
-    """prep + write the HTCondor submit description (one job per plot x chunk), then submit."""
-    import subprocess
+    """Render training-progress plots (loss/lr/accuracy/grad-norm/router/wgan/throughput) from <run_dir>/metrics.csv."""
+    from giant.training.plots import render_metrics

-    from giant.analysis import SubmitConfig, prep, write_submit
-
-    path = prep(
-        rollout_yaml,
-        run_dir,
-        n_chunks=chunks,
-        default_base=Path.cwd() / "analysis_runs",
-        n_energy_bins=n_energy_bins,
-        n_marginal_bins=n_marginal_bins,
-        top_k_pdg=top_k_pdg,
-    )
-    cfg = SubmitConfig(
-        run_dir=path,
-        accounting_group=accounting_group,
-        repo_dir=Path.cwd(),
-        docker_image=docker_image,
-        request_memory_mb=request_memory,
-        remote=remote,
-        n_chunks=chunks,
-    )
-    sub = write_submit(cfg)
-    typer.echo(f"run directory: {path}")
-    typer.echo(f"wrote submit description: {sub}")
-    if dry_run:
-        typer.echo("dry-run: not submitting")
-        return
-    subprocess.run(["condor_submit", str(sub)], check=True)
+    paths = render_metrics(run_dir, out_dir, default_base=Path.cwd() / "analysis_runs")
+    typer.echo(f"rendered {len(paths)} plots -> {paths[0].parent if paths else '(nothing to render)'}")


 if __name__ == "__main__":
@@ -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,3 +1,5 @@
+from __future__ import annotations
+
 import copy
 import difflib
 import hashlib
@@ -10,14 +12,17 @@ from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from enum import Enum
 from pathlib import Path
+from typing import TYPE_CHECKING

-import numpy as np
-import torch
+from giant._migration import V02_FIXED_FACTS, V02_MODEL_KEY_TO_STAGES, reject_legacy_router_expert_sizing
+
+if TYPE_CHECKING:
+    import torch


 class Conditioning(str, Enum):
    """`conditioning.particle.type` / `conditioning.material.type` choices —
-    shared by `giant.cli` and `scripts.dwarf`'s Typer commands so the two
+    shared by `giant.cli` and `giant.tools.dwarf`'s Typer commands so the two
    CLIs can't silently drift apart on the option's valid values (see
    DEFAULT_CONFIG["conditioning"] for what each value means)."""

@@ -47,13 +52,11 @@ CONFIG_VERSION = 3
 # `lambda` is a Python keyword, so dict key "lambda" is always exposed as the
 # field `lambda_weight`.
 #
-# Two sub-blocks — router and n_sec — carry genuinely dynamic keys that don't
-# fit a fixed schema: composed-router `axis{i}_{field}` flags (see
-# giant.model.network._parse_composed_axes) and pipeline.py's runtime-seeded
-# `centers_init`, plus n_sec's `legacy_owner` (injected only by
-# _migrate_legacy_model_config for v0.2 checkpoints). Both dataclasses carry
-# an `extra: dict` catch-all so these keys round-trip losslessly without
-# becoming named fields that would leak into every new run's config.toml.
+# The router sub-block carries genuinely dynamic keys that don't fit a fixed schema:
+# composed-router `axis{i}_{field}` flags (see giant.model.network._parse_composed_axes)
+# and pipeline.py's runtime-seeded `centers_init`. It carries an `extra: dict` catch-all
+# so these keys round-trip losslessly without becoming named fields that would leak into
+# every new run's config.toml.


@dataclass(frozen=True)
@@ -336,6 +339,35 @@ class RouterConfig:
        }


+@dataclass(frozen=True)
+class TrunkConfig:
+    """`stage1_model.trunk`/`stage2_model.trunk`: selects the trunk's expert
+    *body* architecture from `giant.model.trunks.TRUNK_REGISTRY` (default
+    `"resmlp"` — today's only body, `input_proj -> ResBlock stack ->
+    out_proj`). Orthogonal to whether that body is mixed: mixing is still
+    controlled entirely by `router.enabled`/`router.n_experts` on the same
+    stage, unaffected by this block. A future body's own hyperparameters
+    (e.g. a transformer's `n_heads`/`n_layers`) would get their own sibling
+    field here, matching how `flow`/`ddpm`/`wgan` already coexist selected by
+    `generator`.
+
+    `block_conditioning` selects each body's conditioning-injection mechanism
+    from `giant.model.layers.BLOCK_REGISTRY` — `"add"` (default, today's
+    conditional-bias `ResBlock`, bit-identical to pre-gitea-#34 behaviour),
+    `"film"`, or `"adaln"`."""
+
+    type: str = "resmlp"
+    block_conditioning: str = "add"
+
+    @classmethod
+    def from_dict(cls, d: dict | None) -> "TrunkConfig":
+        d = d or {}
+        return cls(type=d.get("type", "resmlp"), block_conditioning=d.get("block_conditioning", "add"))
+
+    def to_dict(self) -> dict:
+        return {"type": self.type, "block_conditioning": self.block_conditioning}
+
+
@dataclass(frozen=True)
 class Stage2RouterConfig(RouterConfig):
    # true: stage 2 shares stage 1's Router module instance, so expert i in
@@ -374,36 +406,63 @@ class Stage2RouterConfig(RouterConfig):
        return {"tie_to_stage1": self.tie_to_stage1, **super().to_dict()}


+# stage2_model.n_sec.sampling choices — single source of truth for both
+# validate_config's train-time check and INFERENCE_OVERRIDES below.
+STOP_SAMPLING_CHOICES = ("greedy", "sample")
+
+# stage2_model.particle_type.other_policy choices — see ParticleTypeConfig's
+# docstring for what each means; only documented there until now, since
+# nothing validated it at train time.
+OTHER_POLICY_CHOICES = ("sample", "modal", "drop")
+
+
@dataclass(frozen=True)
 class NSecConfig:
    # "head": a classifier over {0..k_max} on the condition encoding alone
    # (no diffusion noise), callable independently at inference.
-    # "stop_token": an EOS-style implicit stop — accepted by the schema but
-    # not implemented in v0.3.0 (see validate_config).
+    # "stop_token": an EOS-style per-slot stop head on the autoregressive
+    # secondary decoder (Stage2Autoregressive only — see validate_config),
+    # evaluated against the generated prefix instead of conditioning alone.
+    # Replaces n_sec_head entirely: the two are mutually exclusive, so this
+    # mode builds no n_sec_head and stage2_model.n_sec.lambda instead weights
+    # the stop head's BCE term.
    # "truth": take n_sec from ground truth — standalone stage-2 evaluation
    # only, never for rollout.
    mode: str = "head"
-    lambda_weight: float = 0.1  # dict key "lambda" — cross-entropy weight for the head
-    # Holds "legacy_owner" when injected by _migrate_legacy_model_config
-    # (v0.2 checkpoints only) — not a user-facing config.toml key.
-    extra: dict = field(default_factory=dict)
-
-    @property
-    def legacy_owner(self) -> str | None:
-        return self.extra.get("legacy_owner")
+    lambda_weight: float = 0.1  # dict key "lambda" — cross-entropy/BCE weight for the head
+    # Which stage's module physically owns the n_sec_head weights: "stage2" (default,
+    # fresh v0.3.0 runs — Stage2OneShot/Stage2Autoregressive builds it) or "stage1"
+    # (a migrated v0.2 checkpoint — see network._migrate_legacy_model_config, whose
+    # n_sec head was trained against Stage 1's own ConditionEncoder output and so has
+    # to stay attached there, not just be labeled as such).
+    owner: str = "stage2"
+    # How resolve_n_sec/sample_secondaries_ar turn a count-bearing head's output into an
+    # actual n_sec decision. mode="head": "greedy" is argmax over the classifier logits
+    # (deterministic — the conditional mode, not a sample); "sample" is a categorical draw
+    # from softmax(logits) (a real sample from the learned count distribution). mode=
+    # "stop_token": "greedy" is sigmoid(stop_logit) >= 0.5 per slot (deterministic);
+    # "sample" is a Bernoulli draw at sigmoid(stop_logit) per slot. Renamed from
+    # "stop_sampling" (gitea #86), which is still accepted as a deprecated alias since it
+    # appears in existing checkpoints' model_config.
+    sampling: str = "greedy"

    @classmethod
    def from_dict(cls, d: dict | None) -> "NSecConfig":
        d = d or {}
-        known = {"mode", "lambda"}
        return cls(
            mode=d.get("mode", "head"),
            lambda_weight=d.get("lambda", 0.1),
-            extra={k: v for k, v in d.items() if k not in known},
+            owner=d.get("owner", "stage2"),
+            sampling=d.get("sampling", d.get("stop_sampling", "greedy")),
        )

    def to_dict(self) -> dict:
-        return {"mode": self.mode, "lambda": self.lambda_weight, **self.extra}
+        return {
+            "mode": self.mode,
+            "lambda": self.lambda_weight,
+            "owner": self.owner,
+            "sampling": self.sampling,
+        }


@dataclass(frozen=True)
@@ -420,6 +479,23 @@ class ParticleTypeConfig:
    # at map-build time. "modal": always the most common member. "drop":
    # discard the secondary. Read only under target = "onehot".
    other_policy: str = "sample"
+    # Secondary-species class count under target = "onehot" — independent of
+    # conditioning.particle.emb_dim (see gitea #29: the two used to be
+    # silently the same number). 0 = inherit conditioning.particle.emb_dim,
+    # preserving pre-#29 behavior.
+    n_classes: int = 0
+    # Class-balances the target = "onehot" cross-entropy loss against the
+    # secondary-species long tail (gitea #44: the failure mode motivating the
+    # v0.3.0 pivot was specifically a species collapse — zero photon
+    # secondaries, hallucinated antineutrinos). "none": plain CE (pre-#44
+    # behavior). "inverse_freq": CE weighted by 1/count per class,
+    # normalized to mean 1 across classes so lambda_weight doesn't need
+    # retuning when this is switched on. validate_config requires target =
+    # "onehot" and stage2_model.generator != "wgan" whenever this isn't
+    # "none" — "embedding"/"physical" have no class CE to weight, and the
+    # WGAN stage-2 path feeds its type slice to the critic via a
+    # straight-through Gumbel relaxation instead of a CE loss.
+    class_weighting: str = "none"

    @classmethod
    def from_dict(cls, d: dict | None) -> "ParticleTypeConfig":
@@ -428,10 +504,18 @@ class ParticleTypeConfig:
            target=d.get("target", "onehot"),
            lambda_weight=d.get("lambda", 1.0),
            other_policy=d.get("other_policy", "sample"),
+            n_classes=d.get("n_classes", 0),
+            class_weighting=d.get("class_weighting", "none"),
        )

    def to_dict(self) -> dict:
-        return {"target": self.target, "lambda": self.lambda_weight, "other_policy": self.other_policy}
+        return {
+            "target": self.target,
+            "lambda": self.lambda_weight,
+            "other_policy": self.other_policy,
+            "n_classes": self.n_classes,
+            "class_weighting": self.class_weighting,
+        }


@dataclass(frozen=True)
@@ -477,11 +561,82 @@ class AutoregressiveConfig:
        }


+@dataclass(frozen=True)
+class HeadConfig:
+    """A single classifier head's shape — `n_sec_head`/`type_head` (gitea
+    #36 deduplicated their five identical hand-rolled
+    `Linear -> SiLU -> Linear` definitions into
+    `giant.model.layers.build_mlp_head`, which this config drives).
+    `hidden_ratio=0.5`/`depth=2` are the exact pre-#36 hardcoded values
+    (hidden width = `hidden_dim // 2`, one hidden layer), so omitting a
+    `heads` block — including every migrated v0.2 config — reproduces the
+    old architecture bit-for-bit."""
+
+    hidden_ratio: float = 0.5  # hidden width = round(hidden_dim * hidden_ratio)
+    depth: int = 2  # matches build_mlp_head's depth
+
+    @classmethod
+    def from_dict(cls, d: dict | None) -> "HeadConfig":
+        d = d or {}
+        return cls(hidden_ratio=d.get("hidden_ratio", 0.5), depth=d.get("depth", 2))
+
+    def to_dict(self) -> dict:
+        return {"hidden_ratio": self.hidden_ratio, "depth": self.depth}
+
+
+@dataclass(frozen=True)
+class Stage1HeadsConfig:
+    """Stage 1 only ever owns `n_sec_head`, and only for a migrated v0.2
+    checkpoint (`stage2_model.n_sec.owner = "stage1"`) — see
+    `Stage1Model`'s docstring."""
+
+    n_sec: HeadConfig = field(default_factory=HeadConfig)
+
+    @classmethod
+    def from_dict(cls, d: dict | None) -> "Stage1HeadsConfig":
+        d = d or {}
+        return cls(n_sec=HeadConfig.from_dict(d.get("n_sec")))
+
+    def to_dict(self) -> dict:
+        return {"n_sec": self.n_sec.to_dict()}
+
+
+@dataclass(frozen=True)
+class Stage2HeadsConfig:
+    """`n_sec` and `type` are independently configurable — n_sec accuracy
+    and secondary-species accuracy are separately known weak spots (gitea
+    #36)."""
+
+    n_sec: HeadConfig = field(default_factory=HeadConfig)
+    type: HeadConfig = field(default_factory=HeadConfig)
+
+    @classmethod
+    def from_dict(cls, d: dict | None) -> "Stage2HeadsConfig":
+        d = d or {}
+        return cls(n_sec=HeadConfig.from_dict(d.get("n_sec")), type=HeadConfig.from_dict(d.get("type")))
+
+    def to_dict(self) -> dict:
+        return {"n_sec": self.n_sec.to_dict(), "type": self.type.to_dict()}
+
+
@dataclass(frozen=True)
 class Stage1ModelConfig:
    # false skips building/training stage 1 entirely. The resulting
    # checkpoint holds only stage 2 and cannot be rolled out.
    active: bool = True
+    # Checkpoint .pt to load this stage's weights from before training starts
+    # (its own "model"/"sec_decoder" key, not this run's own resume state) —
+    # "" means start from a fresh init. See `freeze` below for the partial-
+    # retrain use case this exists for (gitea #42).
+    init_from: str = ""
+    # true keeps this stage's weights exactly as loaded from `init_from` —
+    # forward/backward still run every batch (so its loss/grad_norm metrics
+    # stay meaningful, and a WGAN stage's critic still gets a real signal to
+    # report), but its optimizer never steps. Lets a rollout-capable
+    # checkpoint retrain only the *other* stage against a fixed, known-good
+    # one (gitea #42) — `validate_config` requires `init_from` to be set
+    # whenever this is true, unless the run is a `--resume`.
+    freeze: bool = False
    # "flow": conditional flow matching (~10 ODE steps at inference).
    # "ddpm": cosine-schedule diffusion baseline.
    # "wgan": WGAN-GP, single forward pass at inference.
@@ -500,12 +655,16 @@ class Stage1ModelConfig:
    ddpm: DdpmConfig = field(default_factory=DdpmConfig)
    wgan: Stage1WganConfig = field(default_factory=Stage1WganConfig)
    router: RouterConfig = field(default_factory=RouterConfig)
+    trunk: TrunkConfig = field(default_factory=TrunkConfig)
+    heads: Stage1HeadsConfig = field(default_factory=Stage1HeadsConfig)

    @classmethod
    def from_dict(cls, d: dict | None) -> "Stage1ModelConfig":
        d = d or {}
        return cls(
            active=d.get("active", True),
+            init_from=d.get("init_from", ""),
+            freeze=d.get("freeze", False),
            generator=d.get("generator", "flow"),
            hidden_dim=d.get("hidden_dim", 256),
            n_res_blocks=d.get("n_res_blocks", 6),
@@ -515,11 +674,15 @@ class Stage1ModelConfig:
            ddpm=DdpmConfig.from_dict(d.get("ddpm")),
            wgan=Stage1WganConfig.from_dict(d.get("wgan")),
            router=RouterConfig.from_dict(d.get("router")),
+            trunk=TrunkConfig.from_dict(d.get("trunk")),
+            heads=Stage1HeadsConfig.from_dict(d.get("heads")),
        )

    def to_dict(self) -> dict:
        return {
            "active": self.active,
+            "init_from": self.init_from,
+            "freeze": self.freeze,
            "generator": self.generator,
            "hidden_dim": self.hidden_dim,
            "n_res_blocks": self.n_res_blocks,
@@ -529,6 +692,8 @@ class Stage1ModelConfig:
            "ddpm": self.ddpm.to_dict(),
            "wgan": self.wgan.to_dict(),
            "router": self.router.to_dict(),
+            "trunk": self.trunk.to_dict(),
+            "heads": self.heads.to_dict(),
        }


@@ -537,6 +702,9 @@ class Stage2ModelConfig:
    # false trains stage 1 alone. giant rollout must then refuse the
    # checkpoint; giant predict still works.
    active: bool = True
+    # See Stage1ModelConfig.init_from/.freeze — same semantics, this stage.
+    init_from: str = ""
+    freeze: bool = False
    # "one_shot": predict all k_max slots simultaneously with padded slots
    # masked from the loss (v0.2 behaviour).
    # "autoregressive": emit one secondary at a time in descending-energy
@@ -559,6 +727,13 @@ class Stage2ModelConfig:
    # output, closing the train/inference gap at the cost of a sampling pass
    # per batch and a moving target early in training.
    stage1_context: str = "truth"
+    # Ramp for "sampled": P(condition on the ground-truth stage-1 outcome
+    # rather than a fresh sample), linearly interpolated from ctx_p_start
+    # (epoch 0) to ctx_p_end (the final epoch) — the same scheduled-sampling
+    # shape as autoregressive.tf_p_start/tf_p_end, so stage 2 doesn't chase a
+    # wildly moving stage-1 target in early epochs. Unread under "truth".
+    ctx_p_start: float = 1.0
+    ctx_p_end: float = 0.0
    n_sec: NSecConfig = field(default_factory=NSecConfig)
    particle_type: ParticleTypeConfig = field(default_factory=ParticleTypeConfig)
    autoregressive: AutoregressiveConfig = field(default_factory=AutoregressiveConfig)
@@ -566,12 +741,16 @@ class Stage2ModelConfig:
    ddpm: DdpmConfig = field(default_factory=DdpmConfig)
    wgan: Stage2WganConfig = field(default_factory=Stage2WganConfig)
    router: Stage2RouterConfig = field(default_factory=Stage2RouterConfig)
+    trunk: TrunkConfig = field(default_factory=TrunkConfig)
+    heads: Stage2HeadsConfig = field(default_factory=Stage2HeadsConfig)

    @classmethod
    def from_dict(cls, d: dict | None) -> "Stage2ModelConfig":
        d = d or {}
        return cls(
            active=d.get("active", True),
+            init_from=d.get("init_from", ""),
+            freeze=d.get("freeze", False),
            decoder=d.get("decoder", "autoregressive"),
            generator=d.get("generator", "wgan"),
            hidden_dim=d.get("hidden_dim", 256),
@@ -581,6 +760,8 @@ class Stage2ModelConfig:
            k_max=d.get("k_max", 15),
            context_dim=d.get("context_dim", 64),
            stage1_context=d.get("stage1_context", "truth"),
+            ctx_p_start=d.get("ctx_p_start", 1.0),
+            ctx_p_end=d.get("ctx_p_end", 0.0),
            n_sec=NSecConfig.from_dict(d.get("n_sec")),
            particle_type=ParticleTypeConfig.from_dict(d.get("particle_type")),
            autoregressive=AutoregressiveConfig.from_dict(d.get("autoregressive")),
@@ -588,11 +769,15 @@ class Stage2ModelConfig:
            ddpm=DdpmConfig.from_dict(d.get("ddpm")),
            wgan=Stage2WganConfig.from_dict(d.get("wgan")),
            router=Stage2RouterConfig.from_dict(d.get("router")),
+            trunk=TrunkConfig.from_dict(d.get("trunk")),
+            heads=Stage2HeadsConfig.from_dict(d.get("heads")),
        )

    def to_dict(self) -> dict:
        return {
            "active": self.active,
+            "init_from": self.init_from,
+            "freeze": self.freeze,
            "decoder": self.decoder,
            "generator": self.generator,
            "hidden_dim": self.hidden_dim,
@@ -602,6 +787,8 @@ class Stage2ModelConfig:
            "k_max": self.k_max,
            "context_dim": self.context_dim,
            "stage1_context": self.stage1_context,
+            "ctx_p_start": self.ctx_p_start,
+            "ctx_p_end": self.ctx_p_end,
            "n_sec": self.n_sec.to_dict(),
            "particle_type": self.particle_type.to_dict(),
            "autoregressive": self.autoregressive.to_dict(),
@@ -609,6 +796,8 @@ class Stage2ModelConfig:
            "ddpm": self.ddpm.to_dict(),
            "wgan": self.wgan.to_dict(),
            "router": self.router.to_dict(),
+            "trunk": self.trunk.to_dict(),
+            "heads": self.heads.to_dict(),
        }


@@ -642,6 +831,12 @@ class TrainConfig:
    # thousands of steps. Per-epoch metrics (the metrics.csv row) always log
    # in full.
    wandb_log_every: int = 50
+    # Training-step autocast dtype: "fp32" (default, no autocast) or "bf16".
+    # No "fp16" — GradScaler and the double-backward in
+    # giant.model.wgan.gradient_penalty don't mix well, and bf16 alone covers
+    # every training GPU in the fleet (Ampere and newer). See
+    # giant.training.amp.resolve_autocast (gitea #47).
+    precision: str = "fp32"

    @classmethod
    def from_dict(cls, d: dict | None) -> "TrainConfig":
@@ -663,6 +858,7 @@ class TrainConfig:
            wandb_project=d.get("wandb_project", "giant"),
            wandb_run_name=d.get("wandb_run_name", ""),
            wandb_log_every=d.get("wandb_log_every", 50),
+            precision=d.get("precision", "fp32"),
        )

    def to_dict(self) -> dict:
@@ -683,6 +879,7 @@ class TrainConfig:
            "wandb_project": self.wandb_project,
            "wandb_run_name": self.wandb_run_name,
            "wandb_log_every": self.wandb_log_every,
+            "precision": self.precision,
        }


@@ -719,6 +916,25 @@ class GiantConfig:
 DEFAULT_CONFIG: dict = GiantConfig().to_dict()


+def leaf_paths(node: dict, prefix: str = "") -> list[str]:
+    """Every dotted leaf path in a DEFAULT_CONFIG-shaped dict, e.g.
+    "stage1_model.router.n_experts". `[meta]` (run provenance, no schema
+    counterpart) is skipped at the top level, matching `validate_config_keys`.
+    Shared by `tests/test_config_consumed_keys.py` (the static per-identifier
+    audit) and `giant.model.summary` (the runtime per-config audit, gitea
+    #46) so both walk the exact same tree."""
+    paths = []
+    for key, value in node.items():
+        if prefix == "" and key == "meta":
+            continue
+        path = f"{prefix}.{key}" if prefix else key
+        if isinstance(value, dict):
+            paths.extend(leaf_paths(value, path))
+        else:
+            paths.append(path)
+    return paths
+
+
 def git_hash() -> str:
    try:
        return subprocess.check_output(["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL).decode().strip()
@@ -727,6 +943,8 @@ def git_hash() -> str:


 def auto_device() -> torch.device:
+    import torch
+
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
@@ -772,6 +990,8 @@ def estimate_batch_size(
    inference (e.g. `predict`), which uses a much lower per-sample memory
    calibration since there's no backward graph or optimizer state.
    """
+    import torch
+
    if device.type != "cuda":
        raise ValueError(f"--batch-size auto is only supported on cuda devices, got {device.type!r}")
    device_index = device.index if device.index is not None else torch.cuda.current_device()
@@ -875,6 +1095,27 @@ def _set_path(d: dict, dotted: str, value) -> None:
    cur[parts[-1]] = value


+def _pop_path(d: dict, dotted: str) -> None:
+    """Remove a dotted path from a nested dict, if present. No-op if any
+    component along the path is missing."""
+    parts = dotted.split(".")
+    cur = d
+    for part in parts[:-1]:
+        if not isinstance(cur, dict) or part not in cur:
+            return
+        cur = cur[part]
+    if isinstance(cur, dict):
+        cur.pop(parts[-1], None)
+
+
+# Config keys renamed within v0.3 itself (not part of the v0.2->v0.3 migration
+# above) — normalized by migrate_config so a config.toml still using an older
+# v0.3 key name keeps passing validate_config_keys.
+_RENAMED_KEYS = {
+    "stage2_model.n_sec.stop_sampling": "stage2_model.n_sec.sampling",  # gitea #86
+}
+
+
 def _deep_merge(base: dict, override: dict) -> dict:
    """Recursively merge `override` onto a copy of `base`.

@@ -893,6 +1134,72 @@ def _deep_merge(base: dict, override: dict) -> dict:
    return result


+@dataclass(frozen=True)
+class InferenceOverride:
+    """One dotted `model_config` path that `giant.checkpoint_io.load_for_inference`
+    is allowed to change on an already-trained checkpoint, without retraining.
+
+    A path only belongs here if it affects neither module construction/tensor
+    shapes nor the data preprocessing the normalizers/vocab maps were fit
+    under — see the module docstring on `giant.model.summary` for the class
+    of key this targets (`_fingerprint`'s "plain scalar attribute" leaves),
+    and `giant.checkpoint_io.apply_config_overrides` for where this is used.
+    """
+
+    why: str
+    choices: tuple[str, ...] | None = None
+    minimum: float | None = None
+    numeric: bool = False  # int/float leaf (vs. str, the default)
+
+    def check(self, path: str, value: object) -> None:
+        if self.choices is not None:
+            if value not in self.choices:
+                raise ValueError(f"{path} = {value!r} — must be one of {self.choices}")
+            return
+        if self.numeric:
+            if isinstance(value, bool) or not isinstance(value, (int, float)):
+                raise ValueError(f"{path} = {value!r} — must be a number")
+            if self.minimum is not None and value < self.minimum:
+                raise ValueError(f"{path} = {value!r} — must be >= {self.minimum}")
+
+
+# Inference-safe dotted `model_config` paths — the allowlist gitea #87 asked
+# for, so a typo or a shape-bearing key (e.g. "stage1_model.hidden_dim")
+# raises a clear CheckpointCompatibilityError instead of surfacing as an
+# opaque load_state_dict size mismatch later. Extend this table, not a
+# per-call bypass, when a new inference-only key needs the capability.
+INFERENCE_OVERRIDES: dict[str, InferenceOverride] = {
+    "stage2_model.n_sec.sampling": InferenceOverride(
+        why="giant.sample's n_sec head/stop-token sampling reads this at sample time only (gitea #86)",
+        choices=STOP_SAMPLING_CHOICES,
+    ),
+    "stage1_model.ddpm.n_steps": InferenceOverride(
+        why="giant.model.schedule.CosineSchedule's step count, resolved at sample time",
+        numeric=True,
+        minimum=1,
+    ),
+    "stage2_model.ddpm.n_steps": InferenceOverride(
+        why="giant.model.schedule.CosineSchedule's step count, resolved at sample time",
+        numeric=True,
+        minimum=1,
+    ),
+    "stage2_model.particle_type.other_policy": InferenceOverride(
+        why="giant.rollout resolves an 'other'-bucket secondary's PDG code with this at rollout time",
+        choices=OTHER_POLICY_CHOICES,
+    ),
+    "stage1_model.router.temperature": InferenceOverride(
+        why="giant.model.routers.EnergyRouter.temperature, a plain constructor attribute",
+        numeric=True,
+        minimum=1e-6,
+    ),
+    "stage2_model.router.temperature": InferenceOverride(
+        why="giant.model.routers.EnergyRouter.temperature, a plain constructor attribute",
+        numeric=True,
+        minimum=1e-6,
+    ),
+}
+
+
@dataclass(frozen=True)
 class FlagSpec:
    """One CLI flag's mapping into the config-overrides tree.
@@ -936,6 +1243,7 @@ FLAG_SPECS: tuple[FlagSpec, ...] = (
    FlagSpec("wandb_project", ("train.wandb_project",)),
    FlagSpec("wandb_run_name", ("train.wandb_run_name",)),
    FlagSpec("wandb_log_every", ("train.wandb_log_every",)),
+    FlagSpec("precision", ("train.precision",)),
    # --hidden-dim/--n-blocks/--dropout are stage-1-only backward-compat
    # shorthands (they predate stage2_model having its own flags);
    # --stage1-* wins when both are given.
@@ -981,6 +1289,20 @@ FLAG_SPECS: tuple[FlagSpec, ...] = (
    FlagSpec("critic_lr", ("stage1_model.wgan.critic_lr", "stage2_model.wgan.critic_lr"), precedence=0),
    FlagSpec("stage1_critic_lr", ("stage1_model.wgan.critic_lr",), precedence=1),
    FlagSpec("stage2_critic_lr", ("stage2_model.wgan.critic_lr",), precedence=1),
+    # Critic sizing: stage-scoped only, no shared alias — this is an
+    # architectural per-stage knob like hidden_dim/n_res_blocks above, not a
+    # shared training hyperparameter like the wgan knobs above it.
+    FlagSpec("stage1_critic_hidden_dim", ("stage1_model.wgan.critic_hidden_dim",)),
+    FlagSpec("stage1_critic_n_res_blocks", ("stage1_model.wgan.critic_n_res_blocks",)),
+    FlagSpec("stage2_critic_hidden_dim", ("stage2_model.wgan.critic_hidden_dim",)),
+    FlagSpec("stage2_critic_n_res_blocks", ("stage2_model.wgan.critic_n_res_blocks",)),
+    # Partial-retrain (gitea #42): stage-scoped only, no shared alias — a
+    # shared "freeze both stages from the same file" flag has no sensible
+    # meaning (a checkpoint has one set of weights per stage).
+    FlagSpec("stage1_init_from", ("stage1_model.init_from",)),
+    FlagSpec("stage1_freeze", ("stage1_model.freeze",)),
+    FlagSpec("stage2_init_from", ("stage2_model.init_from",)),
+    FlagSpec("stage2_freeze", ("stage2_model.freeze",)),
 )


@@ -1024,14 +1346,6 @@ _V02_TRAIN_PASSTHROUGH = (
    "wandb_log_every",
 )

-# v0.2 model.hidden_dim/n_blocks/dropout applied identically to both stages
-# (there was only ever one trunk shape) -> copied to both stage{1,2}_model.
-_V02_MODEL_TO_BOTH_STAGES = (
-    ("hidden_dim", "hidden_dim"),
-    ("n_blocks", "n_res_blocks"),
-    ("dropout", "dropout"),
-)
-
 # v0.2 train.{n_critic,gp_weight,critic_lr} applied identically to both
 # stages' wgan sub-table (there was only ever one wgan objective, shared).
 _V02_TRAIN_TO_BOTH_STAGES_WGAN = (
@@ -1064,11 +1378,21 @@ def migrate_config(cfg: dict) -> dict:
    (which additionally carries n_sec_head ownership and needs
    `network.build_models`'s cooperation) is a separate migration surface,
    deferred to the network.py refactor.
-    """
-    if _get_path(cfg, "meta.config_version") == CONFIG_VERSION:
-        return copy.deepcopy(cfg)

+    Independently of the v0.2/v0.3 branch below, `_RENAMED_KEYS` normalizes
+    keys renamed within v0.3 itself (e.g. `stop_sampling` -> `sampling`,
+    gitea #86) so a config.toml written against an older v0.3 key name still
+    passes `validate_config_keys`.
+    """
    cfg = copy.deepcopy(cfg)
+    for old_path, new_path in _RENAMED_KEYS.items():
+        if _get_path(cfg, old_path) is not None and _get_path(cfg, new_path) is None:
+            _set_path(cfg, new_path, _get_path(cfg, old_path))
+        _pop_path(cfg, old_path)
+
+    if _get_path(cfg, "meta.config_version") == CONFIG_VERSION:
+        return cfg
+
    old_train = cfg.pop("train", {})
    old_model = cfg.pop("model", {})
    old_router = dict(old_model.pop("router", {}))
@@ -1091,7 +1415,7 @@ def migrate_config(cfg: dict) -> dict:
            _set_path(new, f"stage1_model.wgan.{new_key}", old_train[old_key])
            _set_path(new, f"stage2_model.wgan.{new_key}", old_train[old_key])

-    for old_key, new_key in _V02_MODEL_TO_BOTH_STAGES:
+    for old_key, new_key in V02_MODEL_KEY_TO_STAGES:
        if old_key in old_model:
            _set_path(new, f"stage1_model.{new_key}", old_model[old_key])
            _set_path(new, f"stage2_model.{new_key}", old_model[old_key])
@@ -1108,18 +1432,7 @@ def migrate_config(cfg: dict) -> dict:
        _set_path(new, "stage2_model.k_max", old_model["k_max"])

    if old_router:
-        expert_hidden_dim = old_router.pop("expert_hidden_dim", 0)
-        expert_n_blocks = old_router.pop("expert_n_blocks", 0)
-        if expert_hidden_dim or expert_n_blocks:
-            raise ValueError(
-                "v0.2 config sets model.router.expert_hidden_dim/"
-                f"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 config's routed experts "
-                "have a different width/depth than the monolith and its "
-                "checkpoint can only be loaded by v0.2 code."
-            )
+        reject_legacy_router_expert_sizing(old_router, source="v0.2 config's model.router")
        _set_path(new, "stage1_model.router", dict(old_router))
        stage2_router = dict(old_router)
        stage2_router["tie_to_stage1"] = False
@@ -1127,21 +1440,9 @@ def migrate_config(cfg: dict) -> dict:

    # v0.2 architectural facts with no corresponding config key at all —
    # always set once we've determined we're migrating a v0.2 dict,
-    # independent of what the file did/didn't specify. NOTE: n_layers here
-    # (2) differs from the v0.3 *default* (1) — this is not a typo, see the
-    # docstring above.
-    _set_path(new, "conditioning.out_dim", 128)
-    _set_path(new, "conditioning.particle.n_layers", 2)
-    _set_path(new, "conditioning.material.n_layers", 2)
-    _set_path(new, "stage1_model.active", True)
-    _set_path(new, "stage1_model.flow.time_dim", 64)
-    _set_path(new, "stage1_model.ddpm.time_dim", 64)
-    _set_path(new, "stage2_model.active", True)
-    _set_path(new, "stage2_model.flow.time_dim", 64)
-    _set_path(new, "stage2_model.ddpm.time_dim", 64)
-    _set_path(new, "stage2_model.context_dim", 64)
-    _set_path(new, "stage2_model.decoder", "one_shot")
-    _set_path(new, "stage2_model.particle_type.target", "physical")
+    # independent of what the file did/didn't specify (see giant._migration).
+    for path, value in V02_FIXED_FACTS.items():
+        _set_path(new, path, value)

    new_meta = dict(cfg.pop("meta", {}))
    new_meta["config_version"] = CONFIG_VERSION
@@ -1244,7 +1545,7 @@ def merge_cli_overrides(
    return cfg


-def validate_config(cfg: dict) -> None:
+def validate_config(cfg: dict, *, resume: bool = False) -> None:
    """Cross-block validation the per-block schema can't express on its own.

    Raises ValueError with a clear message on the first violation found. Call
@@ -1252,6 +1553,10 @@ def validate_config(cfg: dict) -> None:
    these checks need to see across blocks, so they don't belong in
    `migrate_config` (which only ever sees one dict's own keys) or in any
    single block's defaults.
+
+    `resume=True` (only `giant train --resume` passes this) relaxes the
+    `stage{1,2}_model.freeze` -> `.init_from` requirement below: a resumed
+    frozen stage's weights come from the resume checkpoint, not `init_from`.
    """
    particle_type = _get_path(cfg, "conditioning.particle.type")

@@ -1264,7 +1569,34 @@ def validate_config(cfg: dict) -> None:
            f"{particle_type!r})"
        )

+    class_weighting = _get_path(cfg, "stage2_model.particle_type.class_weighting")
+    if class_weighting not in ("none", "inverse_freq"):
+        raise ValueError(
+            f"stage2_model.particle_type.class_weighting = {class_weighting!r} — must be 'none' or 'inverse_freq'"
+        )
+    if class_weighting != "none" and pt_target != "onehot":
+        raise ValueError(
+            "stage2_model.particle_type.class_weighting != 'none' requires "
+            f"stage2_model.particle_type.target = 'onehot' (there is no class "
+            f"cross-entropy to weight under target = {pt_target!r})"
+        )
+    if class_weighting != "none" and _get_path(cfg, "stage2_model.generator") == "wgan":
+        raise ValueError(
+            "stage2_model.particle_type.class_weighting != 'none' is "
+            "incompatible with stage2_model.generator = 'wgan' — that path "
+            "feeds the type slice to the critic via a straight-through "
+            "Gumbel relaxation instead of a class cross-entropy, so there is "
+            "nothing to weight"
+        )
+
    for stage_name in ("stage1_model", "stage2_model"):
+        if _get_path(cfg, f"{stage_name}.freeze") and not _get_path(cfg, f"{stage_name}.init_from") and not resume:
+            raise ValueError(
+                f"{stage_name}.freeze = true requires {stage_name}.init_from "
+                "to be set (or --resume) — freezing a randomly-initialized "
+                "model is almost certainly a mistake"
+            )
+
        router = _get_path(cfg, f"{stage_name}.router") or {}
        if router.get("enabled") and router.get("type") in ("pdg", "process") and particle_type == "physical":
            raise ValueError(
@@ -1283,12 +1615,54 @@ def validate_config(cfg: dict) -> None:
        )

    if _get_path(cfg, "stage2_model.n_sec.mode") == "stop_token":
+        if _get_path(cfg, "stage2_model.decoder") != "autoregressive":
+            raise ValueError(
+                "stage2_model.n_sec.mode = 'stop_token' requires "
+                "stage2_model.decoder = 'autoregressive' — there is no "
+                "per-token loop to stop under 'one_shot'"
+            )
+        if _get_path(cfg, "stage2_model.n_sec.owner") != "stage2":
+            raise ValueError(
+                "stage2_model.n_sec.mode = 'stop_token' requires "
+                "stage2_model.n_sec.owner = 'stage2' — a migrated v0.2 "
+                "checkpoint's stage-1 n_sec_head has no per-token "
+                "conditioning to hang an EOS decision off"
+            )
+
+    n_sec_sampling = _get_path(cfg, "stage2_model.n_sec.sampling")
+    if n_sec_sampling not in STOP_SAMPLING_CHOICES:
+        raise ValueError(f"stage2_model.n_sec.sampling = {n_sec_sampling!r} — must be 'greedy' or 'sample'")
+
+    precision = _get_path(cfg, "train.precision")
+    if precision not in ("fp32", "bf16"):
        raise ValueError(
-            "stage2_model.n_sec.mode = 'stop_token' is accepted by the schema "
-            "but not implemented in v0.3.0 — use 'head' (default) or 'truth' "
-            "(standalone stage-2 evaluation only, never for rollout)"
+            f"train.precision = {precision!r} — must be 'fp32' or 'bf16' "
+            "('fp16' is not supported: see giant.training.amp.resolve_autocast)"
        )

+    stage1_context = _get_path(cfg, "stage2_model.stage1_context")
+    if stage1_context not in ("truth", "sampled"):
+        raise ValueError(f"stage2_model.stage1_context = {stage1_context!r} — must be 'truth' or 'sampled'")
+    if stage1_context == "sampled":
+        if not (_get_path(cfg, "stage1_model.active") and _get_path(cfg, "stage2_model.active")):
+            raise ValueError(
+                "stage2_model.stage1_context = 'sampled' requires both "
+                "stage1_model.active and stage2_model.active = true — there is "
+                "no stage-1 model to sample from in a stage-2-only run"
+            )
+        ctx_p_start = _get_path(cfg, "stage2_model.ctx_p_start")
+        ctx_p_end = _get_path(cfg, "stage2_model.ctx_p_end")
+        for name, value in (("ctx_p_start", ctx_p_start), ("ctx_p_end", ctx_p_end)):
+            if not (0.0 <= value <= 1.0):
+                raise ValueError(f"stage2_model.{name} = {value} — must be in [0, 1]")
+        if ctx_p_start == 1.0 and ctx_p_end == 1.0:
+            raise ValueError(
+                "stage2_model.stage1_context = 'sampled' with ctx_p_start = "
+                "ctx_p_end = 1.0 always conditions on the ground truth — "
+                "identical to 'truth' but silently so; use 'truth' instead or "
+                "lower ctx_p_end"
+            )
+
    if (
        _get_path(cfg, "stage2_model.n_sec.mode") == "truth"
        and _get_path(cfg, "stage1_model.active")
@@ -1305,9 +1679,20 @@ def validate_config(cfg: dict) -> None:
        )

    if _get_path(cfg, "stage2_model.decoder") == "autoregressive":
+        order = _get_path(cfg, "stage2_model.autoregressive.order")
+        if order != "energy_desc":
+            raise ValueError(
+                f"stage2_model.autoregressive.order = {order!r} — must be "
+                "'energy_desc' (the only implemented ordering; see "
+                "AutoregressiveConfig.order's docstring)"
+            )
+        from giant.model.history import HISTORY_REGISTRY
+
        history = _get_path(cfg, "stage2_model.autoregressive.history")
-        if history not in ("markov", "attention"):
-            raise ValueError(f"stage2_model.autoregressive.history = {history!r} — must be 'markov' or 'attention'")
+        if history not in HISTORY_REGISTRY:
+            raise ValueError(
+                f"stage2_model.autoregressive.history = {history!r} — must be one of {sorted(HISTORY_REGISTRY)}"
+            )
        teacher_forcing = _get_path(cfg, "stage2_model.autoregressive.teacher_forcing")
        if teacher_forcing not in ("always", "scheduled", "never"):
            raise ValueError(
@@ -1400,6 +1785,10 @@ _OUT_DIR_NAME_CANDIDATES = [
        "particle_type_target",
        _path_candidate("stage2_model.particle_type.target", "pt-"),
    ),
+    ("stage1_trunk_type", _path_candidate("stage1_model.trunk.type", "s1t-")),
+    ("stage2_trunk_type", _path_candidate("stage2_model.trunk.type", "s2t-")),
+    ("stage1_block_cond", _path_candidate("stage1_model.trunk.block_conditioning", "s1bc-")),
+    ("stage2_block_cond", _path_candidate("stage2_model.trunk.block_conditioning", "s2bc-")),
    ("stage1_router", _router_candidate("stage1_model", "s1")),
    ("stage2_router", _router_candidate("stage2_model", "s2")),
    (
@@ -1436,6 +1825,8 @@ _OUT_DIR_NAME_CANDIDATES = [
    ),
    ("particle_conditioning", _conditioning_candidate("particle", "c")),
    ("material_conditioning", _conditioning_candidate("material", "m")),
+    ("stage1_freeze", _path_candidate("stage1_model.freeze", "s1frozen", formatter=lambda _: "")),
+    ("stage2_freeze", _path_candidate("stage2_model.freeze", "s2frozen", formatter=lambda _: "")),
    ("stage1_hidden_dim", _path_candidate("stage1_model.hidden_dim", "h")),
    ("stage2_hidden_dim", _path_candidate("stage2_model.hidden_dim", "s2h")),
    ("stage1_n_res_blocks", _path_candidate("stage1_model.n_res_blocks", "b")),
@@ -1497,7 +1888,23 @@ def resolve_default_out_dir(cfg: dict, base: Path = Path("checkpoints")) -> Path
    return out_dir


+def epoch_seed(seed: int, epoch: int) -> int:
+    """Per-epoch derivative of the run seed.
+
+    Reseeding the global RNGs from this at the top of every epoch makes epoch
+    *k* draw the same noise whether it runs inside one long `giant train` or
+    as its own resumed job in a per-epoch workflow chain
+    (`giant/workflow/tasks.py:TrainEpochTask`) — without it, a fresh process
+    would restart the stream at epoch 1's state. Mirrors what
+    `StreamingStepsDataset.set_epoch` does for the batch order.
+    """
+    return (int(seed) * 1_000_003 + int(epoch)) % (2**32)
+
+
 def seed_everything(seed: int) -> None:
+    import numpy as np
+    import torch
+
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
@@ -1551,6 +1958,8 @@ def build_run_meta(
    n_val_events: int,
    n_train_steps: int,
 ) -> dict:
+    import torch
+
    return {
        "config_version": CONFIG_VERSION,
        "git_hash": git_hash(),
@@ -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,24 +71,7 @@ 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"
-        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` — it sets the padded
@@ -82,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)}
@@ -102,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] = []
@@ -129,17 +164,7 @@ class StreamingStepsDataset(IterableDataset):
                    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,
@@ -155,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:
                    (
@@ -217,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]

@@ -226,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
@@ -256,24 +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).
+    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]:
@@ -287,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


@@ -299,6 +307,11 @@ 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:
@@ -315,8 +328,8 @@ def build_topn_map_from_files(files: list[Path], column: str, n_classes: int, ca
    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:
@@ -347,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,
@@ -131,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()},
    }


@@ -139,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()},
    )


@@ -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)
@@ -683,27 +693,31 @@ def decode_secondaries(
    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 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.
+    `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
@@ -720,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:
@@ -729,12 +743,10 @@ 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
@@ -747,14 +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],
@@ -772,49 +836,16 @@ def build_cond_features(
    `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: 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

@@ -822,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.

@@ -831,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
@@ -842,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."
@@ -860,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],
@@ -877,37 +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` — 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
@@ -919,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
@@ -950,29 +988,8 @@ 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)
-
-    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)
+    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)

    n_sec_raw = data["n_sec"].astype(np.int64)  # (N,) unclamped, for the valid-slot mask

@@ -1039,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:
@@ -1054,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,
    )
@@ -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,235 @@
+"""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,
+                n_sec_sampling=s2_spec.n_sec.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(),
+                n_sec_sampling=s2_spec.n_sec.sampling,
+            )
+
+    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,785 @@
+"""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,
+        n_sec_sampling: str = "greedy",
+    ) -> 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.n_sec_sampling = n_sec_sampling
+        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,
+        n_sec_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.n_sec_sampling = n_sec_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
@@ -0,0 +1,328 @@
+"""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 INFERENCE_OVERRIDES, _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)
+    overridable: 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
+    .n_sec_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)
+
+    overridable = sorted(p for p in in_scope if p in INFERENCE_OVERRIDES)
+
+    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),
+        overridable=overridable,
+    )
+
+
+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.overridable:
+        lines.append("")
+        lines.append("inference-overridable without retraining (giant predict/rollout --set):")
+        for path in summary.overridable:
+            lines.append(f"  {path}  ({INFERENCE_OVERRIDES[path].why})")
+
+    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(
@@ -140,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: 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`)
@@ -150,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"`
@@ -31,7 +31,7 @@ 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.model.network import build_models, build_critics, resolve_type_n_classes
 from giant.training import train as run_training


@@ -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
@@ -175,29 +176,42 @@ def run_setup_stage(
                cache.proc_maps[n_experts] = proc_map

    # Top-N-plus-other maps for onehot conditioning/type axes.
-    # 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.
+    # 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 = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type")).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, {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":
@@ -248,7 +262,7 @@ def run_setup_stage(
                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(
+                feats = build_features(
                    chunk_tr,
                    pdg_map,
                    mat_map,
@@ -257,8 +271,19 @@ def run_setup_stage(
                    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:
@@ -292,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,
@@ -333,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"]
@@ -381,8 +407,8 @@ def run_train_job(
    # (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:
@@ -409,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,
@@ -486,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,
@@ -117,7 +117,7 @@ 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]:
@@ -141,7 +141,7 @@ def decode_secondary_identity(
    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()
@@ -158,15 +158,15 @@ def decode_secondary_identity(
    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,
@@ -444,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,
@@ -469,14 +470,17 @@ def rollout(
    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` (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), and
-    `pdg_topn_map`/`other_policy` are additionally read under
+    `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). `seed` seeds the `other_policy = "sample"` draw only
-    (torch/numpy sampling itself is seeded by the caller, same as today).
+    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 embedding-distance
    diagnostic across the whole run — see `L1DistCollector`. Only populated
@@ -487,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 "
@@ -536,6 +545,7 @@ def rollout(
                    material_conditioning,
                    pdg_topn_map,
                    mat_topn_map,
+                    sec_type_topn_map,
                    other_policy,
                    rng,
                    stage1_ddpm_steps,
@@ -574,6 +584,7 @@ def _step_chunk(
    material_conditioning,
    pdg_topn_map,
    mat_topn_map,
+    sec_type_topn_map,
    other_policy,
    rng,
    stage1_ddpm_steps,
@@ -670,7 +681,6 @@ def _step_chunk(
    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)

    # --- Secondaries ---
    # No snapping for "physical"/history-facing state elsewhere in the
@@ -680,7 +690,13 @@ 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_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,
@@ -690,7 +706,7 @@ def _step_chunk(
        tr["pre_dir"],
        sec_phys_norm,
        pdg_map,
-        pdg_topn_map,
+        sec_type_topn_map,
        other_policy,
        rng,
    )
@@ -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


@@ -131,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(
@@ -229,14 +229,14 @@ 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: one token at a time, in
-    descending-energy slot order, `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
+    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.
@@ -244,7 +244,30 @@ def sample_secondaries_ar(
    A `{flow,ddpm}` token costs `steps` ODE substeps; `wgan` costs one pass —
    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.n_sec_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)
@@ -281,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

@@ -294,6 +317,16 @@ 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)
@@ -301,7 +334,27 @@ def sample_secondaries_ar(
        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.n_sec_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,
@@ -362,7 +415,8 @@ def sample_secondaries_ar(
        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


@@ -383,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)
@@ -397,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
@@ -406,10 +460,16 @@ 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":
+    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)

@@ -421,20 +481,38 @@ 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. 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."""
+    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.
+
+    `n_sec.mode = "head"` resolves the classifier logits per
+    `sec_decoder.n_sec_sampling`: "greedy" (default) takes the conditional
+    mode via argmax; "sample" draws a real sample from the learned count
+    distribution via `torch.multinomial` on the softmax — see gitea #86."""
    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"
+            "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)
+    if sec_decoder.n_sec_sampling == "sample":
+        return torch.multinomial(logits.softmax(dim=-1), 1).squeeze(-1)
    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
@@ -579,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)
 # ---------------------------------------------------------------------------


@@ -10,9 +10,9 @@ 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 ("one flow checkpoint and one WGAN checkpoint").
@@ -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$")

@@ -1,10 +1,12 @@
 """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.
 """

+from __future__ import annotations
+
 import os
 from enum import Enum
 from pathlib import Path
@@ -14,20 +16,15 @@ import typer
 from typing_extensions import Annotated

 from giant.config import Conditioning
-from scripts.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
+
+# DATA_DEFAULT/SCAN_DIR_DEFAULT are Typer option defaults (evaluated at
+# decoration time below), so that one name has to stay eager — the module
+# itself is stdlib-only, so it costs nothing. Every other giant.tools.*
+# import here is deferred into the one command body that uses it, since
+# several (steps_to_parquet: uproot/awkward/polars; warm_setup_cache:
+# giant.pipeline -> torch; geometry_oracle: pandas) are expensive and
+# `dwarf --help`/tab-completion shouldn't pay for all of them upfront.
+from giant.tools.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT

 app = typer.Typer(no_args_is_help=True)

@@ -121,6 +118,9 @@ def convert(
    ] = None,
 ) -> None:
    """Convert ROOT Steps tree(s) to Parquet."""
+    from giant.tools.steps_to_parquet import convert_steps_to_parquet
+    from giant.tools.steps_to_parquet_parallel import run_parallel_job
+
    if jobs < 1:
        typer.echo("error: --jobs must be >= 1", err=True)
        raise typer.Exit(1)
@@ -183,6 +183,8 @@ def migrate(
    ] = False,
 ) -> None:
    """One-time migration into the versioned raw/processed/pools/derived layout."""
+    from giant.tools.migrate_geant_steps import run_migration
+
    run_migration(str(root), execute=execute, copy=copy)


@@ -207,6 +209,8 @@ def bump_gen(
    root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
    """Cut a new raw generation."""
+    from giant.tools.bump_dataset_version import run_bump_gen
+
    run_bump_gen(
        kind=kind,
        reason=reason,
@@ -240,6 +244,8 @@ def bump_schema(
    root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
    """Cut a new schema within a gen."""
+    from giant.tools.bump_dataset_version import run_bump_schema
+
    run_bump_schema(
        kind=kind,
        gen=gen,
@@ -257,6 +263,8 @@ def status(
    root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
    """List existing gens/schemas per kind."""
+    from giant.tools.bump_dataset_version import run_status
+
    run_status(str(root))


@@ -281,6 +289,8 @@ def update_manifest(
    ] = False,
 ) -> None:
    """Repoint manifest(s) to a new gen and/or schema, verifying all target files exist."""
+    from giant.tools.bump_dataset_version import run_update_manifest
+
    run_update_manifest([str(m) for m in manifests], schema=schema, execute=execute, gen=gen)


@@ -311,6 +321,8 @@ def create_manifest(
    ] = False,
 ) -> None:
    """Create a new manifest from a list of parquet files."""
+    from giant.tools.bump_dataset_version import run_create_manifest
+
    run_create_manifest(
        [str(f) for f in files],
        execute=execute,
@@ -358,6 +370,8 @@ def make_root(
    ] = False,
 ) -> None:
    """Generate new ROOT shards via a minicalosim executable."""
+    from giant.tools.create_root_files import run_make_root
+
    _warn_if_exceeds_shared_quota(jobs, "--jobs")
    run_make_root(
        executable=executable,
@@ -423,6 +437,8 @@ def build_geometry_oracle(
    ] = 2000,
 ) -> None:
    """Fit a position -> (material, layer_id) oracle for `giant rollout`."""
+    from giant.tools.geometry_oracle import run_build_geometry_oracle
+
    run_build_geometry_oracle(
        data=data,
        out=out,
@@ -442,43 +458,65 @@ def warm_cache(
        Path,
        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,
-        typer.Option("--seed", "-s", help="Must match the `giant train` run(s) to warm for"),
-    ] = 0,
+        Optional[int],
+        typer.Option(
+            "--seed",
+            "-s",
+            help="Must match the `giant train` run(s) to warm for. Not allowed together with --config",
+        ),
+    ] = 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 "
-            "(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,
-    router_type: Annotated[str, typer.Option("--router-type", help="Router implementation name")] = "energy",
-    n_experts: Annotated[int, typer.Option("--n-experts", help="Number of routed experts")] = 4,
+    ] = None,
+    router_type: Annotated[
+        Optional[str],
+        typer.Option("--router-type", help="Router implementation name. Not allowed together with --config"),
+    ] = None,
+    n_experts: Annotated[
+        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"),
@@ -487,17 +525,40 @@ def warm_cache(
    """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.
    """
+    from giant.tools.warm_setup_cache import run_warm_setup_cache
+
+    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,
@@ -514,6 +575,8 @@ def hparam_scan(
    dry_run: Annotated[bool, typer.Option("--dry-run")] = False,
 ) -> None:
    """Grid-scan dropout x n_blocks x hidden_dim via sequential `giant train` runs."""
+    from giant.tools.hparam_scan import run_hparam_scan
+
    run_hparam_scan(data=data, scan_dir=scan_dir, seed=seed, dry_run=dry_run)


@@ -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
@@ -165,11 +166,10 @@ def _time(spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path)
    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,
    )
@@ -191,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+)$")

@@ -82,7 +82,7 @@ def resolve_destination(root_file: Path, dataset_root: Path, schema_override: st
    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(
@@ -127,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
@@ -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.")
@@ -5,7 +5,7 @@ Split out of the former single-module `giant/train.py`. The public surface is
 that tests and tooling construct directly.
 """

-from giant.training.checkpoint import build_checkpoint, load_checkpoint
+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 (
@@ -25,6 +25,7 @@ __all__ = [
    "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)"
+    )
@@ -8,6 +8,8 @@ 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
@@ -46,6 +48,37 @@ def build_checkpoint(
    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
@@ -18,9 +18,10 @@ 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, load_checkpoint
+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,
@@ -109,6 +110,7 @@ def train(
    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,
@@ -131,9 +133,12 @@ def train(
    validate_steps = t.get("validate_steps", 10)
    max_val_batches = t.get("max_val_batches", 0)

-    trainers = build_stage_trainers(cfg, models, critics, device, total_train_batches)
+    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 = {
@@ -142,6 +147,7 @@ def train(
        "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,
    }
@@ -179,6 +185,17 @@ def train(
            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()

@@ -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
@@ -11,7 +11,9 @@ 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


@@ -30,7 +32,7 @@ def _gumbel_tau(step: int, total_steps: int, tau_start: float, tau_end: float) -
 def _type_repr(
    sec_type_idx: torch.Tensor,
    sec_cont: torch.Tensor,
-    particle_type_cfg: dict,
+    particle_type_cfg: ParticleTypeConfig,
    cond_enc: torch.nn.Module,
    emb_dim: int,
 ) -> torch.Tensor:
@@ -45,7 +47,7 @@ def _type_repr(
    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.get("target", "physical")
+    target = particle_type_cfg.target
    if target == "physical":
        return sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM]
    if target == "onehot":
@@ -56,7 +58,7 @@ def _type_repr(
 def _assemble_stage2_ar_target(
    sec_cont: torch.Tensor,
    sec_type_idx: torch.Tensor,
-    particle_type_cfg: dict,
+    particle_type_cfg: ParticleTypeConfig,
    generator: str,
    cond_enc: torch.nn.Module,
    emb_dim: int,
@@ -69,19 +71,20 @@ def _assemble_stage2_ar_target(

    - `target = "physical"`: unchanged from v0.2 — `sec_cont` (stick_logit,
      dir, log_mass, charge) as-is.
-    - `target` in `("onehot", "embedding")` + `generator in ("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")` + `generator == "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` 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.get("target", "physical")
+    target = particle_type_cfg.target
    if target == "physical":
        return sec_cont
    cont = sec_cont[..., :CONT_SLOT_DIM]
-    if generator != "wgan":
+    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)
@@ -90,7 +93,7 @@ def _assemble_stage2_ar_target(
 def _assemble_stage2_real(
    sec_cont: torch.Tensor,
    sec_type_idx: torch.Tensor,
-    particle_type_cfg: dict,
+    particle_type_cfg: ParticleTypeConfig,
    generator: str,
    cond_enc: torch.nn.Module,
    emb_dim: int,
@@ -117,9 +120,17 @@ def _remaining_energy_fraction(fraction: torch.Tensor) -> torch.Tensor:
    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)."""
-    cumprod = torch.cumprod(1.0 - fraction, dim=1)
-    return torch.cat([torch.ones_like(cumprod[:, :1]), cumprod[:, :-1]], dim=1)
+    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:
@@ -137,6 +148,27 @@ def _ar_has_prev(k_max: int, device: torch.device) -> torch.Tensor:
    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
@@ -155,7 +187,7 @@ def _ar_meta(k_max: int, batch: int, device: torch.device, fraction: torch.Tenso
 def _assemble_stage2_ar_inputs(
    sec_cont: torch.Tensor,
    sec_type_idx: torch.Tensor,
-    particle_type_cfg: dict,
+    particle_type_cfg: ParticleTypeConfig,
    cond_enc: torch.nn.Module,
    emb_dim: int,
 ) -> dict[str, torch.Tensor]:
@@ -178,27 +210,43 @@ def _assemble_stage2_ar_inputs(
    return {"history_feat": history_feat, **_ar_meta(K, B, device, fraction)}


-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` (epoch 0) to `p_end` (the final epoch) — standard scheduled
-    sampling (Bengio et al. 2015)."""
-    if mode == "always":
-        return 1.0
-    if mode == "never":
-        return 0.0
+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: dict,
+    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
@@ -211,7 +259,7 @@ def _history_repr_from_ar_sample(
    representation."""
    fraction = torch.sigmoid(sec_cont_pred[..., 0])
    direction = sec_cont_pred[..., 1:4]
-    if particle_type_cfg.get("target", "physical") == "onehot":
+    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:
@@ -227,7 +275,7 @@ def _assemble_stage2_ar_inputs_scheduled(
    sec_cont: torch.Tensor,
    sec_type_idx: torch.Tensor,
    n_sec: torch.Tensor,
-    particle_type_cfg: dict,
+    particle_type_cfg: ParticleTypeConfig,
    cond_enc: torch.nn.Module,
    emb_dim: int,
    p_tf: float,
@@ -16,6 +16,7 @@ adversarial and non-adversarial stages identically.
 import copy
 import math
 from dataclasses import dataclass, field
+from typing import NamedTuple

 import torch
 import torch.nn.functional as F
@@ -23,21 +24,20 @@ import torch.optim as optim

 from giant.config import ParticleTypeConfig, Stage1ModelConfig, Stage2ModelConfig, TrainConfig
 from giant.constants import CONT_SLOT_DIM
-from giant.model.network import Router, stage2_type_dim
-from giant.model.schedule import (
-    CosineSchedule,
-    flow_matching_loss,
-    flow_matching_loss_secondary,
-    flow_matching_loss_secondary_ar,
-)
+from giant.data.dataset import StepBatch
+from giant.model.network import Router, build_objective, resolve_type_n_classes, stage2_type_dim
 from giant.model.wgan import generator_loss, gradient_penalty
+from giant.sample import sample_stage1
+from giant.training.amp import resolve_autocast
 from giant.training.metrics import MetricSpec, stage_metric, train_metric, val_metric
 from giant.training.stage2_inputs import (
    _assemble_stage2_ar_inputs_scheduled,
    _assemble_stage2_ar_target,
+    _ctx_truth_prob,
    _gumbel_tau,
    _relax_onehot_type_slice,
    _stage2_tf_prob,
+    _stop_target_and_mask,
 )


@@ -72,8 +72,43 @@ def _cosine_warmup_lambda(warmup_steps: int, total_steps: int):
    return _lr_lambda


-def _batch_to_device(batch: tuple, device: torch.device) -> tuple:
-    return tuple(t.to(device) for t in batch)
+def _batch_to_device(batch: StepBatch, device: torch.device) -> StepBatch:
+    return type(batch)(*(t.to(device) for t in batch))
+
+
+def _type_class_weight_vector(class_counts: dict[int, int], n_classes: int, scheme: str) -> list[float] | None:
+    """Per-class `F.cross_entropy(weight=...)` vector for the stage-2 type
+    head's `class_weighting` (gitea #44), or `None` under `"none"` (the
+    pre-#44 unweighted-CE behavior — the caller must pass that through as
+    `weight=None`, not a vector of ones, so old runs stay bit-identical).
+
+    `"inverse_freq"`: `1 / count` per class, normalized to mean 1 over
+    `n_classes` so switching this on doesn't rescale the type loss against
+    `particle_type.lambda` / the generator loss it's summed with. A class
+    with zero training examples (fewer distinct species than `n_classes - 1`
+    slots) clamps its count to 1 — its weight is otherwise undefined, and
+    since it never appears in a batch's labels the value is inert anyway.
+
+    Raises if `scheme != "none"` and `class_counts` is empty: that means the
+    `TopNMap` behind this run predates gitea #44 (a stale checkpoint's decode
+    map, or a not-yet-rebuilt setup-cache sidecar) and truly has no
+    frequency information to weight by — silently falling back to uniform
+    weights would look like the feature is active when it isn't.
+    """
+    if scheme == "none":
+        return None
+    if not class_counts:
+        raise ValueError(
+            f"stage2_model.particle_type.class_weighting = {scheme!r} requires "
+            "per-class counts, but this run's sec_type_topn_map has none "
+            "(class_counts={}) — it was built before gitea #44 or loaded "
+            "from a stale setup-cache sidecar/checkpoint; rebuild the setup "
+            "cache (giant train --rebuild-setup-cache) or retrain."
+        )
+    counts = [max(class_counts.get(i, 0), 1) for i in range(n_classes)]
+    inv = [1.0 / c for c in counts]
+    mean_inv = sum(inv) / len(inv)
+    return [w / mean_inv for w in inv]


@dataclass(frozen=True)
@@ -90,13 +125,22 @@ class StageSpec:
    generator: str
    decoder: str = "one_shot"

+    # partial-retrain (gitea #42)
+    init_from: str = ""
+    freeze: bool = False
+
    # loss weights
    lambda_weight: float = 1.0
    n_sec_lambda: float = 0.1
+    n_sec_mode: str = "head"

    # particle-type target (stage 2 only)
    particle_type: ParticleTypeConfig = field(default_factory=ParticleTypeConfig)
-    particle_type_emb_dim: int = 16
+    particle_type_n_classes: int = 16
+    # Resolved by from_config from sec_type_class_counts (dataset-derived,
+    # not itself a cfg value — see _type_class_weight_vector) crossed with
+    # particle_type.class_weighting (gitea #44). None under "none".
+    type_class_weights: list[float] | None = None

    # optimization
    lr: float = 3e-4
@@ -105,6 +149,7 @@ class StageSpec:
    warmup_epochs: int = 0
    epochs: int = 1
    steps_per_epoch: int = 1
+    precision: str = "fp32"

    # routing auxiliaries
    lambda_balance: float = 0.0
@@ -119,6 +164,11 @@ class StageSpec:
    tf_p_end: float = 1.0
    ar_sample_steps: int = 10

+    # stage-1/stage-2 boundary (stage 2 only)
+    stage1_context: str = "truth"
+    ctx_p_start: float = 1.0
+    ctx_p_end: float = 0.0
+
    # generator-specific
    ddpm_n_steps: int = 1000
    n_critic: int = 5
@@ -128,7 +178,18 @@ class StageSpec:
    type_gumbel_tau_end: float = 0.1

    @classmethod
-    def from_config(cls, cfg: dict, name: str, is_stage2: bool, steps_per_epoch: int) -> "StageSpec":
+    def from_config(
+        cls,
+        cfg: dict,
+        name: str,
+        is_stage2: bool,
+        steps_per_epoch: int,
+        sec_type_class_counts: dict[int, int] | None = None,
+    ) -> "StageSpec":
+        """`sec_type_class_counts` is dataset-derived (`sec_type_topn_map.class_counts`,
+        gitea #44), not a `cfg` value — it's the one input to `StageSpec` that
+        doesn't come from `cfg`, kept separate from the "only place that
+        reads `cfg`" invariant below on purpose."""
        t = TrainConfig.from_dict(cfg["train"])
        # n_sec/particle_type/decoder/autoregressive/wgan's gumbel_tau_* are
        # stage-2-only concepts, always read off s2_spec (guarded by
@@ -141,15 +202,24 @@ class StageSpec:
        # stage 1's).
        s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"])
        stage_spec = s2_spec if is_stage2 else Stage1ModelConfig.from_dict(cfg["stage1_model"])
+        particle_type_n_classes = resolve_type_n_classes(
+            s2_spec.particle_type, cfg["conditioning"]["particle"]["emb_dim"]
+        )
        return cls(
            name=name,
            is_stage2=is_stage2,
            generator=stage_spec.generator,
            decoder=s2_spec.decoder if is_stage2 else "one_shot",
+            init_from=stage_spec.init_from,
+            freeze=stage_spec.freeze,
            lambda_weight=stage_spec.lambda_weight,
            n_sec_lambda=s2_spec.n_sec.lambda_weight,
+            n_sec_mode=s2_spec.n_sec.mode,
            particle_type=s2_spec.particle_type,
-            particle_type_emb_dim=cfg["conditioning"]["particle"]["emb_dim"],
+            particle_type_n_classes=particle_type_n_classes,
+            type_class_weights=_type_class_weight_vector(
+                sec_type_class_counts or {}, particle_type_n_classes, s2_spec.particle_type.class_weighting
+            ),
            # train.* keys are all guaranteed by DEFAULT_CONFIG's deep-merge
            # (giant/config.py), so TrainConfig.from_dict never has to fall
            # back to a literal here; the field defaults below exist only
@@ -160,6 +230,7 @@ class StageSpec:
            warmup_epochs=t.warmup_epochs,
            epochs=t.epochs,
            steps_per_epoch=max(steps_per_epoch, 1),
+            precision=t.precision,
            lambda_balance=stage_spec.router.lambda_balance,
            lambda_proc=stage_spec.router.lambda_proc,
            lambda_entropy=stage_spec.router.lambda_entropy,
@@ -168,6 +239,9 @@ class StageSpec:
            teacher_forcing=s2_spec.autoregressive.teacher_forcing if is_stage2 else cls.teacher_forcing,
            tf_p_start=s2_spec.autoregressive.tf_p_start if is_stage2 else cls.tf_p_start,
            tf_p_end=s2_spec.autoregressive.tf_p_end if is_stage2 else cls.tf_p_end,
+            stage1_context=s2_spec.stage1_context if is_stage2 else cls.stage1_context,
+            ctx_p_start=s2_spec.ctx_p_start if is_stage2 else cls.ctx_p_start,
+            ctx_p_end=s2_spec.ctx_p_end if is_stage2 else cls.ctx_p_end,
            # AR self-sampling under scheduled/never teacher forcing reuses
            # train.validate_steps as its flow-matching ODE step count — no
            # dedicated config key for this (the autoregressive config lists
@@ -185,14 +259,16 @@ class StageSpec:
 class StageTrainer:
    """One active stage's optimizer(s), EMA, and per-batch step.

-    Reads only the shared batch tuple `(cond_cont, cond_cat, x1_s1, n_sec,
-    sec_cont, proc_idx, sec_type_idx)` — stage 2 always conditions on the
-    ground-truth `x1_s1` (`stage2_model.stage1_context = "truth"`,
-    stage-level teacher forcing; `"sampled"` is not implemented), so stage
-    trainers never need each other's output at train time. This means
-    "stage-2-only training is a cheap ablation, not new plumbing" falls out
-    for free: a trainer only exists for active stages, and inactive stages
-    are simply never constructed.
+    Reads only the shared `StepBatch` (`giant.data.dataset`) by default — stage
+    2 conditions on the ground-truth `x1_s1` (`stage2_model.stage1_context =
+    "truth"`, stage-level teacher forcing), so "stage-2-only training is a
+    cheap ablation, not new plumbing" falls out for free: a trainer only
+    exists for active stages, and inactive stages are simply never
+    constructed. `stage2_model.stage1_context = "sampled"` is the one
+    exception — `build_stage_trainers` wires the stage-2 trainer to the
+    stage-1 one via `attach_stage1` so it can draw a real stage-1 sample
+    (`giant.sample.sample_stage1`) instead, scheduled by `ctx_p_start`/
+    `ctx_p_end` (see `_stage1_context`).

    Grad-norm clipping is per-stage here — v0.2's single shared optimizer
    clipped both stages' gradients jointly; splitting per stage is a small,
@@ -231,13 +307,20 @@ class StageTrainer:
        self.is_stage2 = spec.is_stage2
        self.generator = spec.generator
        self.decoder = spec.decoder
+        self.frozen = spec.freeze
        self.device = device
        self.model = model.to(device)
        self.router = _stage_router(self.model)
        self._modules = (self.model, *extra_modules)

-        self.particle_type_cfg = spec.particle_type.to_dict()
-        self.particle_type_emb_dim = spec.particle_type_emb_dim
+        # Resolved once (not re-derived every step) — see
+        # giant.training.amp.resolve_autocast (gitea #47).
+        self._autocast_device_type, self._autocast_dtype, self._autocast_enabled = resolve_autocast(
+            spec.precision, device
+        )
+
+        self.particle_type_cfg = spec.particle_type
+        self.particle_type_n_classes = spec.particle_type_n_classes
        self.ema_decay = spec.ema_decay

        self.ema_model: torch.nn.Module | None = None
@@ -246,6 +329,18 @@ class StageTrainer:
            for p in self.ema_model.parameters():
                p.requires_grad_(False)

+        #: Set by `attach_stage1` when `stage2_model.stage1_context =
+        #: "sampled"` — the stage-1 `StageTrainer` this (stage-2) trainer
+        #: draws its context sample from. `None` for stage 1 itself, and for
+        #: stage 2 under "truth".
+        self.stage1_source: "StageTrainer | None" = None
+
+    def attach_stage1(self, stage1_trainer: "StageTrainer") -> None:
+        """Wires this (stage-2) trainer to the stage-1 trainer it should
+        sample from under `stage2_model.stage1_context = "sampled"` — see
+        `build_stage_trainers`."""
+        self.stage1_source = stage1_trainer
+
    # --- schedule -------------------------------------------------------

    def _init_lr_schedule(self, optimizer: optim.Optimizer, warmup_steps: int, total_steps: int) -> None:
@@ -255,10 +350,10 @@ class StageTrainer:

    # --- per-batch (subclass responsibility) ----------------------------

-    def step(self, batch: tuple, device: torch.device, global_step: int) -> dict:
+    def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict:
        raise NotImplementedError

-    def val_loss(self, batch: tuple, device: torch.device) -> dict:
+    def val_loss(self, batch: StepBatch, device: torch.device) -> dict:
        raise NotImplementedError

    # --- reporting hooks ------------------------------------------------
@@ -289,6 +384,58 @@ class StageTrainer:
        for module in self._modules:
            module.eval()

+    # --- stage-1/stage-2 boundary (shared by both trainer subclasses) ---
+
+    def _stage1_context(
+        self,
+        x1_s1: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        epoch: int | None,
+    ) -> torch.Tensor:
+        """The stage-1 outcome stage 2 conditions on this batch.
+
+        `epoch=None` means "always ground truth" regardless of
+        `spec.stage1_context` — the same val-loss convention `_ar_inputs`
+        uses, so validation stays a stable, non-stochastic comparison.
+        Otherwise, under `stage1_context = "sampled"`, each example
+        independently uses the ground truth with probability `p_truth`
+        (`_ctx_truth_prob`, ramped by `ctx_p_start`/`ctx_p_end`) and a fresh
+        `giant.sample.sample_stage1` draw from `stage1_source.sampling_model()`
+        otherwise — a real sampling pass, not a cheap proxy, matching
+        `_assemble_stage2_ar_inputs_scheduled`'s precedent for the equivalent
+        in-stage-2 self-sample. Mixed per example (not per-dimension): a row
+        is either the real ground-truth 9D vector or a real sample, never an
+        elementwise blend of the two.
+        """
+        x1_s1 = x1_s1.detach()
+        if self.stage1_source is None or epoch is None:
+            return x1_s1
+        p_truth = _ctx_truth_prob(
+            self.spec.stage1_context,
+            self.spec.ctx_p_start,
+            self.spec.ctx_p_end,
+            epoch,
+            self.spec.epochs,
+        )
+        if p_truth >= 1.0:
+            return x1_s1
+
+        stage1_model = self.stage1_source.sampling_model()
+        was_training = stage1_model.training
+        sampled, _ = sample_stage1(
+            stage1_model,
+            cond_cont,
+            cond_cat,
+            steps=self.spec.ar_sample_steps,
+            ddpm_steps=self.stage1_source.spec.ddpm_n_steps,
+        )
+        if was_training:
+            stage1_model.train()
+
+        use_truth = torch.rand(x1_s1.size(0), 1, device=x1_s1.device) < p_truth
+        return torch.where(use_truth, x1_s1, sampled).detach()
+
    # --- stage-2 secondary assembly (shared by both trainer subclasses) ---

    def _ar_inputs(
@@ -328,7 +475,7 @@ class StageTrainer:
            n_sec,
            self.particle_type_cfg,
            self.model.cond_enc,
-            self.particle_type_emb_dim,
+            self.particle_type_n_classes,
            p_tf,
            self.spec.ar_sample_steps,
        )
@@ -354,7 +501,7 @@ class StageTrainer:
            self.particle_type_cfg,
            generator,
            self.model.cond_enc,
-            self.particle_type_emb_dim,
+            self.particle_type_n_classes,
        )
        return target.flatten(1) if flatten else target

@@ -377,10 +524,9 @@ class StageTrainer:
        stage1-vs-stage2 `predict_n_sec` signature split, shared by the
        non-adversarial and WGAN trainers.

-        Gated on `n_sec_head is None`, not on `n_sec.mode`: a future
-        `mode="stop_token"` model (currently rejected in
-        `validate_config`) carries no head and would train its EOS signal in
-        the generator/AR loss path instead, so this correctly stays zero.
+        Gated on `n_sec_head is None`, not on `n_sec.mode`: a `mode =
+        "stop_token"` model carries no head at all (see `_stop_loss` for its
+        EOS signal instead), so this correctly stays zero for it.
        """
        if self.model.n_sec_head is None:
            zero = torch.zeros((), device=device)
@@ -394,14 +540,77 @@ class StageTrainer:
        nsec_acc = (logits.argmax(dim=-1) == n_sec).float().mean()
        return l_nsec, nsec_acc

-    @staticmethod
-    def _step_optimizer(optimizer: optim.Optimizer, loss: torch.Tensor, params: list) -> float:
+    def _stop_loss(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_ctx: torch.Tensor,
+        n_sec: torch.Tensor,
+        device: torch.device,
+        ar_inputs: dict[str, torch.Tensor] | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """`(l_stop, stop_acc)` for `n_sec.mode = "stop_token"`'s per-slot EOS
+        head (`Stage2Autoregressive.predict_stop`) — zeros when this stage
+        owns no `stop_head` (every other `n_sec.mode`), the same gating
+        convention `_n_sec_loss` uses for `n_sec_head`. The two heads are
+        mutually exclusive (`giant.model.builders`), so exactly one of
+        `_n_sec_loss`/`_stop_loss` is ever non-zero for a given stage.
+
+        Masked BCE against `_stop_target_and_mask`'s per-slot target — one
+        slot wider than `sec_mask` (the stop slot itself, `k == n_sec`, needs
+        supervision even though it holds no real secondary)."""
+        stop_head = getattr(self.model, "stop_head", None)
+        if stop_head is None:
+            zero = torch.zeros((), device=device)
+            return zero, zero
+        assert ar_inputs is not None
+        logits = self.model.predict_stop(
+            cond_cont,
+            cond_cat,
+            stage1_ctx,
+            ar_inputs["history_feat"],
+            ar_inputs["has_prev"],
+            ar_inputs["remaining_frac"],
+            ar_inputs["slot_idx"],
+        )
+        target, mask = _stop_target_and_mask(n_sec, logits.size(1), device)
+        mask_f = mask.float()
+        denom = mask_f.sum().clamp(min=1)
+        bce = F.binary_cross_entropy_with_logits(logits, target, reduction="none")
+        l_stop = (bce * mask_f).sum() / denom
+        stop_acc = (((logits >= 0).float() == target).float() * mask_f).sum() / denom
+        return l_stop, stop_acc
+
+    def _autocast(self) -> torch.autocast:
+        """The training-step autocast region (`train.precision`, gitea #47).
+
+        Only wraps forward/loss computation — `backward()`/`optimizer.step()`
+        stay outside, and `val_loss` never calls this at all, so validation
+        (and the best-checkpoint selection it drives) stays precision-
+        independent and comparable against every fp32-only run recorded so
+        far. `enabled=False` under `precision = "fp32"` (the default) makes
+        this a true no-op, so callers never need to branch on precision
+        themselves."""
+        return torch.autocast(
+            self._autocast_device_type,
+            dtype=self._autocast_dtype,
+            enabled=self._autocast_enabled,
+        )
+
+    def _step_optimizer(self, optimizer: optim.Optimizer, loss: torch.Tensor, params: list) -> float:
        """`zero_grad -> backward -> clip_grad_norm_(1.0) -> step`, returning
-        the pre-clip grad norm. The one place the grad-clip constant lives."""
+        the pre-clip grad norm. The one place the grad-clip constant lives.
+
+        `self.frozen` (`stage{1,2}_model.freeze`, gitea #42) skips only the
+        final `optimizer.step()` — backward/clip still run so loss/grad_norm
+        stay meaningful to watch, but the stage's weights (and, for a WGAN
+        stage, its critic's — this same method is both trainers' single
+        optimizer-step choke point) never move."""
        optimizer.zero_grad()
        loss.backward()
        grad_norm = torch.nn.utils.clip_grad_norm_(params, 1.0)
-        optimizer.step()
+        if not self.frozen:
+            optimizer.step()
        return grad_norm.item()

    def _extra_state(self) -> dict:
@@ -447,20 +656,29 @@ class FlowDDPMStageTrainer(StageTrainer):
    """flow or ddpm generator for a single stage."""

    def __init__(self, spec: StageSpec, model: torch.nn.Module, device: torch.device) -> None:
-        if spec.is_stage2 and spec.generator not in ("flow",):
+        objective = build_objective(spec.generator, n_steps=spec.ddpm_n_steps)
+        if spec.is_stage2 and not objective.supports_stage2_decoder:
            raise NotImplementedError(
                f"stage2_model.generator={spec.generator!r} is accepted by the "
                "schema but not implemented in v0.3.0 for stage 2 (only "
                "'flow' and 'wgan' have a stage-2 secondary-decoder loss)"
            )
        super().__init__(spec, model, device)
-        self.particle_type_lambda = self.particle_type_cfg.get("lambda", 1.0)
+        self.objective = objective
+        self.particle_type_lambda = self.particle_type_cfg.lambda_weight
        # Width of the type slice actually folded into x1_s2 by _sec_target,
-        # under this trainer's generator (flow/ddpm only — see the
-        # NotImplementedError above): "physical" keeps it folded in
-        # (PARTICLE_PHYS_DIM wide, unchanged from v0.2); "onehot"/"embedding"
-        # pull it out into model.type_head instead (0 here).
-        self._flow_type_dim = None if self.particle_type_cfg.get("target", "physical") == "physical" else 0
+        # under this trainer's objective (flow/ddpm only — see the
+        # NotImplementedError above, neither folds the type slice): "physical"
+        # keeps it folded in (PARTICLE_PHYS_DIM wide, unchanged from v0.2);
+        # "onehot"/"embedding" pull it out into model.type_head instead (0
+        # here).
+        self._flow_type_dim = None if self.particle_type_cfg.target == "physical" else 0
+        # gitea #44: None under class_weighting = "none" (the default),
+        # matching F.cross_entropy's own unweighted default — a real tensor
+        # only materializes when the config asked for one.
+        self.type_class_weights = (
+            None if spec.type_class_weights is None else torch.tensor(spec.type_class_weights, device=device)
+        )

        self.params = list(self.model.parameters())
        self.optimizer = optim.AdamW(self.params, lr=spec.lr, weight_decay=spec.weight_decay)
@@ -469,7 +687,7 @@ class FlowDDPMStageTrainer(StageTrainer):
            warmup_steps=spec.warmup_epochs * spec.steps_per_epoch,
            total_steps=max(spec.epochs * spec.steps_per_epoch, 1),
        )
-        self.ddpm_schedule = CosineSchedule(T=spec.ddpm_n_steps).to(device) if spec.generator == "ddpm" else None
+        self.ddpm_schedule = self.objective.build_schedule(spec.ddpm_n_steps, device)

        self.train_metrics = [
            train_metric(key)
@@ -477,10 +695,12 @@ class FlowDDPMStageTrainer(StageTrainer):
                "loss",
                "loss_gen",
                "loss_nsec",
+                "loss_stop",
                "loss_balance",
                "loss_proc",
                "loss_entropy",
                "nsec_acc",
+                "stop_acc",
                "loss_type",
                "type_acc",
                "grad_norm",
@@ -493,6 +713,8 @@ class FlowDDPMStageTrainer(StageTrainer):
                "loss_gen",
                "loss_nsec",
                "nsec_acc",
+                "loss_stop",
+                "stop_acc",
                "loss_type",
                "type_acc",
            )
@@ -501,26 +723,9 @@ class FlowDDPMStageTrainer(StageTrainer):

    def _generator_loss(self, cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=None):
        if not self.is_stage2:
-            if self.generator == "flow":
-                return flow_matching_loss(self.model, x1_s1, cond_cont, cond_cat)
-            assert self.ddpm_schedule is not None
-            return self.ddpm_schedule.loss(self.model, x1_s1, cond_cont, cond_cat)
-        if self.decoder == "autoregressive":
-            assert ar_inputs is not None
-            return flow_matching_loss_secondary_ar(
-                self.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=self._flow_type_dim,
-            )
-        return flow_matching_loss_secondary(
+            return self.objective.stage1_loss(self.model, x1_s1, cond_cont, cond_cat, schedule=self.ddpm_schedule)
+        assert self.decoder != "autoregressive" or ar_inputs is not None
+        return self.objective.stage2_loss(
            self.model,
            x1_s2,
            cond_cont,
@@ -528,6 +733,7 @@ class FlowDDPMStageTrainer(StageTrainer):
            stage1_ctx,
            sec_mask,
            type_dim=self._flow_type_dim,
+            ar_inputs=ar_inputs,
        )

    def _type_loss(
@@ -565,9 +771,13 @@ class FlowDDPMStageTrainer(StageTrainer):
            type_out = self.model.predict_type(cond_cont, cond_cat, stage1_ctx)
        mask = sec_mask.float()
        denom = mask.sum().clamp(min=1)
-        if self.particle_type_cfg.get("target") == "onehot":
-            ce = F.cross_entropy(type_out.transpose(1, 2), sec_type_idx, reduction="none")
+        if self.particle_type_cfg.target == "onehot":
+            ce = F.cross_entropy(
+                type_out.transpose(1, 2), sec_type_idx, weight=self.type_class_weights, reduction="none"
+            )
            l_type = (ce * mask).sum() / denom
+            # Unweighted, deliberately — type_acc is a diagnostic of raw
+            # per-slot correctness, not the (possibly class-weighted) loss.
            type_acc = ((type_out.argmax(-1) == sec_type_idx).float() * mask).sum() / denom
        else:  # "embedding"
            target_vec = self.model.cond_enc.pdg_emb(sec_type_idx).detach()
@@ -575,11 +785,12 @@ class FlowDDPMStageTrainer(StageTrainer):
            l_type = (se * mask).sum() / denom
        return l_type, type_acc

-    def _compute(self, batch: tuple, device: torch.device, epoch: int | None = None) -> dict:
+    def _compute(self, batch: StepBatch, device: torch.device, epoch: int | None = None) -> dict:
        """`epoch=None` (the `val_loss` path) always uses full teacher
-        forcing (`p_tf=1.0`) regardless of `spec.teacher_forcing` — validation
-        should stay a stable, non-stochastic ground-truth comparison; only
-        the training `step` path schedules `p_tf` by epoch."""
+        forcing (`p_tf=1.0`) and the ground-truth stage-1 context, regardless
+        of `spec.teacher_forcing`/`spec.stage1_context` — validation should
+        stay a stable, non-stochastic ground-truth comparison; only the
+        training `step` path schedules `p_tf`/`p_truth` by epoch."""
        (
            cond_cont,
            cond_cat,
@@ -590,7 +801,7 @@ class FlowDDPMStageTrainer(StageTrainer):
            sec_type_idx,
        ) = _batch_to_device(batch, device)
        sec_mask = self._sec_mask(n_sec, sec_cont.size(1), device)
-        stage1_ctx = x1_s1.detach()
+        stage1_ctx = self._stage1_context(x1_s1, cond_cont, cond_cat, epoch)

        x1_s2 = None
        ar_inputs = None
@@ -602,6 +813,7 @@ class FlowDDPMStageTrainer(StageTrainer):

        l_gen = self._generator_loss(cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=ar_inputs)
        l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device)
+        l_stop, stop_acc = self._stop_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device, ar_inputs)

        l_type, type_acc = self._type_loss(
            cond_cont,
@@ -615,11 +827,18 @@ class FlowDDPMStageTrainer(StageTrainer):

        l_balance = l_proc = l_entropy = torch.zeros((), device=device)
        if self.router is not None:
-            l_balance = self.router.balance_loss(cond_cont, cond_cat)
-            l_proc = self.router.classify_loss(cond_cont, cond_cat, proc_idx)
-            l_entropy = self.router.entropy_loss(cond_cont, cond_cat)
+            if self.spec.lambda_balance > 0:
+                l_balance = self.router.balance_loss(cond_cont, cond_cat)
+            if self.spec.lambda_proc > 0:
+                l_proc = self.router.classify_loss(cond_cont, cond_cat, proc_idx)
+            if self.spec.lambda_entropy > 0:
+                l_entropy = self.router.entropy_loss(cond_cont, cond_cat)

-        total = self.spec.lambda_weight * l_gen + self.spec.n_sec_lambda * l_nsec + self.particle_type_lambda * l_type
+        total = (
+            self.spec.lambda_weight * l_gen
+            + self.spec.n_sec_lambda * (l_nsec + l_stop)
+            + self.particle_type_lambda * l_type
+        )
        if self.spec.lambda_balance > 0:
            total = total + self.spec.lambda_balance * l_balance
        if self.spec.lambda_proc > 0:
@@ -631,15 +850,17 @@ class FlowDDPMStageTrainer(StageTrainer):
            "loss": total,
            "loss_gen": l_gen,
            "loss_nsec": l_nsec,
+            "loss_stop": l_stop,
            "loss_type": l_type,
            "type_acc": type_acc,
            "loss_balance": l_balance,
            "loss_proc": l_proc,
            "loss_entropy": l_entropy,
            "nsec_acc": nsec_acc,
+            "stop_acc": stop_acc,
        }

-    def step(self, batch: tuple, device: torch.device, global_step: int) -> dict:
+    def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict:
        if self.router is not None:
            self.router.gumbel_tau = _gumbel_tau(
                global_step,
@@ -648,10 +869,12 @@ class FlowDDPMStageTrainer(StageTrainer):
                self.spec.gumbel_tau_end,
            )
        epoch = global_step // self.spec.steps_per_epoch
-        out = self._compute(batch, device, epoch=epoch)
+        with self._autocast():
+            out = self._compute(batch, device, epoch=epoch)
        grad_norm = self._step_optimizer(self.optimizer, out["loss"], self.params)
-        self.lr_sched.step()
-        if self.ema_model is not None:
+        if not self.frozen:
+            self.lr_sched.step()
+        if self.ema_model is not None and not self.frozen:
            _update_ema(self.ema_model, self.model, self.ema_decay)
        stats = {key: value.item() for key, value in out.items()}
        stats["grad_norm"] = grad_norm
@@ -659,7 +882,7 @@ class FlowDDPMStageTrainer(StageTrainer):
        return stats

    @torch.no_grad()
-    def val_loss(self, batch: tuple, device: torch.device) -> dict:
+    def val_loss(self, batch: StepBatch, device: torch.device) -> dict:
        return {key: value.item() for key, value in self._compute(batch, device).items()}

    # --- reporting ------------------------------------------------------
@@ -674,6 +897,16 @@ class FlowDDPMStageTrainer(StageTrainer):
        return val_means.get("loss", 0.0)


+class _Stage2RealFakeBatch(NamedTuple):
+    """Subset of `StepBatch` that `_stage2_real_and_fake` needs."""
+
+    cond_cont: torch.Tensor
+    cond_cat: torch.Tensor
+    n_sec: torch.Tensor
+    sec_cont: torch.Tensor
+    sec_type_idx: torch.Tensor
+
+
 class WGANStageTrainer(StageTrainer):
    """WGAN-GP generator+critic for a single stage (see giant/model/wgan.py).

@@ -726,10 +959,12 @@ class WGANStageTrainer(StageTrainer):
            "gp_loss",
            "loss_nsec",
            "nsec_acc",
+            "loss_stop",
+            "stop_acc",
            "grad_norm_d",
            "grad_norm_g",
        ]
-        if self.is_stage2 and self.particle_type_cfg.get("target") == "onehot":
+        if self.is_stage2 and self.particle_type_cfg.target == "onehot":
            # Differentiability instrumentation — only meaningful when the
            # type slice is a straight-through Gumbel relaxation.
            train_keys += ["grad_norm_type_slice", "grad_norm_cont_slice"]
@@ -737,15 +972,19 @@ class WGANStageTrainer(StageTrainer):
        self.val_metrics = []
        self.stage_metrics = [stage_metric("lr"), stage_metric("critic_lr")]

-    def _stage2_real_and_fake(self, batch_tensors, stage1_ctx, global_step, device):
-        """Build `(real, fake_raw, mask, critic_fn)` for stage 2, covering
-        both decoders and all three particle-type targets. `fake_raw` still
-        needs the caller's straight-through relaxation under
+    def _stage2_real_and_fake(self, batch_tensors: _Stage2RealFakeBatch, stage1_ctx, global_step, device):
+        """Build `(real, fake_raw, mask, critic_fn, ar_inputs)` for stage 2,
+        covering both decoders and all three particle-type targets. `fake_raw`
+        still needs the caller's straight-through relaxation under
        `particle_type.target = "onehot"`, and neither tensor is masked-and-
-        multiplied on the fake side yet."""
+        multiplied on the fake side yet. `ar_inputs` is `None` under
+        `decoder = "one_shot"`; under `"autoregressive"` it's the same dict
+        `_ar_inputs` built to condition `self.model` above — returned so the
+        caller's `_stop_loss` reuses it instead of paying for a second
+        (possibly self-sampling) `_ar_inputs` call."""
        cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx = batch_tensors
        B = cond_cont.size(0)
-        type_dim = stage2_type_dim(self.particle_type_cfg, self.particle_type_emb_dim)
+        type_dim = stage2_type_dim(self.particle_type_cfg, self.particle_type_n_classes)
        slot_width = CONT_SLOT_DIM + type_dim
        k_max = sec_cont.size(1)

@@ -755,29 +994,30 @@ class WGANStageTrainer(StageTrainer):
        def critic_fn(x):
            return self.critic(x, cond_cont, cond_cat, stage1_ctx)

+        ar_inputs = None
        if self.decoder == "autoregressive":
            epoch = global_step // self.spec.steps_per_epoch
-            ar = self._ar_inputs(cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch)
-            real = self._sec_target(sec_cont, sec_type_idx, "wgan", flatten=False).reshape(B, -1) * mask
+            ar_inputs = self._ar_inputs(cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch)
+            real = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=False).reshape(B, -1) * mask
            z = torch.randn(B, k_max, self.model.noise_dim, device=device)
            fake_raw = self.model(
                z,
                cond_cont,
                cond_cat,
                stage1_ctx,
-                ar["history_feat"],
-                ar["has_prev"],
-                ar["remaining_frac"],
-                ar["slot_idx"],
+                ar_inputs["history_feat"],
+                ar_inputs["has_prev"],
+                ar_inputs["remaining_frac"],
+                ar_inputs["slot_idx"],
            ).reshape(B, -1)
        else:
-            real = self._sec_target(sec_cont, sec_type_idx, "wgan", flatten=True) * mask
+            real = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=True) * mask
            z = torch.randn(B, self.model.noise_dim, device=device)
            fake_raw = self.model(z, cond_cont, cond_cat, stage1_ctx)

-        return real, fake_raw, mask, critic_fn
+        return real, fake_raw, mask, critic_fn, ar_inputs

-    def step(self, batch: tuple, device: torch.device, global_step: int) -> dict:
+    def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict:
        (
            cond_cont,
            cond_cat,
@@ -788,52 +1028,58 @@ class WGANStageTrainer(StageTrainer):
            sec_type_idx,
        ) = _batch_to_device(batch, device)
        B = cond_cont.size(0)
-        stage1_ctx = x1_s1.detach()
+        epoch = global_step // self.spec.steps_per_epoch
+        stage1_ctx = self._stage1_context(x1_s1, cond_cont, cond_cat, epoch)
        grad_probe: dict[str, float] = {}

-        if not self.is_stage2:
-            real = x1_s1
+        ar_inputs = None
+        with self._autocast():
+            if not self.is_stage2:
+                real = x1_s1

-            def critic_fn(x):
-                return self.critic(x, cond_cont, cond_cat)
+                def critic_fn(x):
+                    return self.critic(x, cond_cont, cond_cat)

-            z = torch.randn(B, self.model.noise_dim, device=device)
-            fake = self.model(z, cond_cont, cond_cat)
-            mask = None
-        else:
-            real, fake_raw, mask, critic_fn = self._stage2_real_and_fake(
-                (cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx),
-                stage1_ctx,
-                global_step,
-                device,
-            )
-            if self.particle_type_cfg.get("target", "physical") == "onehot":
-                # Straight-through Gumbel-softmax relaxation of the type
-                # slice only — the critic must see a hard one-hot forward
-                # (matching what "real" data looks like) while gradient
-                # still flows smoothly to the generator. grad_probe captures
-                # the gradient-magnitude instrumentation — see
-                # _relax_onehot_type_slice's docstring.
-                tau = _gumbel_tau(
+                z = torch.randn(B, self.model.noise_dim, device=device)
+                fake = self.model(z, cond_cont, cond_cat)
+                mask = None
+            else:
+                real, fake_raw, mask, critic_fn, ar_inputs = self._stage2_real_and_fake(
+                    _Stage2RealFakeBatch(cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx),
+                    stage1_ctx,
                    global_step,
-                    self.total_steps,
-                    self.spec.type_gumbel_tau_start,
-                    self.spec.type_gumbel_tau_end,
+                    device,
                )
-                fake_raw = _relax_onehot_type_slice(
-                    fake_raw,
-                    sec_cont.size(1),
-                    CONT_SLOT_DIM,
-                    stage2_type_dim(self.particle_type_cfg, self.particle_type_emb_dim),
-                    tau,
-                    grad_probe=grad_probe,
-                )
-            fake = fake_raw * mask
+                if self.particle_type_cfg.target == "onehot":
+                    # Straight-through Gumbel-softmax relaxation of the type
+                    # slice only — the critic must see a hard one-hot forward
+                    # (matching what "real" data looks like) while gradient
+                    # still flows smoothly to the generator. grad_probe captures
+                    # the gradient-magnitude instrumentation — see
+                    # _relax_onehot_type_slice's docstring.
+                    tau = _gumbel_tau(
+                        global_step,
+                        self.total_steps,
+                        self.spec.type_gumbel_tau_start,
+                        self.spec.type_gumbel_tau_end,
+                    )
+                    fake_raw = _relax_onehot_type_slice(
+                        fake_raw,
+                        sec_cont.size(1),
+                        CONT_SLOT_DIM,
+                        stage2_type_dim(self.particle_type_cfg, self.particle_type_n_classes),
+                        tau,
+                        grad_probe=grad_probe,
+                    )
+                fake = fake_raw * mask

-        # --- critic step (every batch) ---
-        fake_detached = fake.detach()
-        real_score = critic_fn(real)
-        fake_score = critic_fn(fake_detached)
+            # --- critic step (every batch) ---
+            fake_detached = fake.detach()
+            real_score = critic_fn(real)
+            fake_score = critic_fn(fake_detached)
+
+        # gradient_penalty forces its own fp32 region internally (see its
+        # docstring) regardless of the ambient autocast above.
        gp = gradient_penalty(critic_fn, real, fake_detached, mask=mask)
        d_loss = fake_score.mean() - real_score.mean() + self.gp_weight * gp
        wasserstein = (real_score.mean() - fake_score.mean()).detach()
@@ -842,27 +1088,31 @@ class WGANStageTrainer(StageTrainer):

        # --- generator (+ n_sec) step ---
        did_g_step = global_step % self.n_critic == 0
-        l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device)
+        with self._autocast():
+            l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device)
+            l_stop, stop_acc = self._stop_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device, ar_inputs)

-        # On a non-generator-step batch with no n_sec_head on this stage
-        # (n_sec now defaults to stage 2), there's nothing for
+        # On a non-generator-step batch with no n_sec_head/stop_head on this
+        # stage (n_sec now defaults to stage 2), there's nothing for
        # the generator optimizer to do this batch — g_loss would otherwise
        # be a graph-less zero tensor, which .backward() rejects outright.
-        skip_g_step = not did_g_step and self.model.n_sec_head is None
+        skip_g_step = not did_g_step and self.model.n_sec_head is None and self.model.stop_head is None
        if did_g_step:
-            g_loss_adv = generator_loss(critic_fn, fake)
-            g_loss = self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * l_nsec
+            with self._autocast():
+                g_loss_adv = generator_loss(critic_fn, fake)
+            g_loss = self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * (l_nsec + l_stop)
        else:
            g_loss_adv = torch.zeros((), device=device)
-            g_loss = self.spec.n_sec_lambda * l_nsec
+            g_loss = self.spec.n_sec_lambda * (l_nsec + l_stop)
        if skip_g_step:
            grad_norm_g = 0.0
        else:
            grad_norm_g = self._step_optimizer(self.optimizer, g_loss, self.g_params)

        if did_g_step:
-            self.lr_sched.step()
-            if self.ema_model is not None:
+            if not self.frozen:
+                self.lr_sched.step()
+            if self.ema_model is not None and not self.frozen:
                _update_ema(self.ema_model, self.model, self.ema_decay)

        return {
@@ -872,6 +1122,8 @@ class WGANStageTrainer(StageTrainer):
            "gp_loss": gp.item(),
            "loss_nsec": l_nsec.item(),
            "nsec_acc": nsec_acc.item(),
+            "loss_stop": l_stop.item(),
+            "stop_acc": stop_acc.item(),
            "did_g_step": did_g_step,
            "grad_norm": grad_norm_d + grad_norm_g,
            "grad_norm_d": grad_norm_d,
@@ -932,16 +1184,28 @@ def build_stage_trainers(
    critics: dict[str, torch.nn.Module | None],
    device: torch.device,
    total_train_batches: int,
+    sec_type_class_counts: dict[int, int] | None = None,
 ) -> dict[str, StageTrainer]:
    """One trainer per active stage — `models[name] is None` means that stage
-    is `active = false` and is simply never constructed."""
+    is `active = false` and is simply never constructed.
+
+    `stage2_model.stage1_context = "sampled"` additionally wires the
+    stage-2 trainer to the stage-1 one (`StageTrainer.attach_stage1`) so it
+    can draw a real stage-1 sample instead of only ever seeing the
+    ground-truth stage-1 outcome — `validate_config` already guarantees both
+    stages are active whenever that config value is set.
+
+    `sec_type_class_counts` (`sec_type_topn_map.class_counts`, gitea #44) is
+    the one dataset-derived input `StageSpec.from_config` needs beyond `cfg`
+    — `None`/absent whenever `stage2_model.particle_type.class_weighting =
+    "none"` (the default), which never reads it."""
    trainers: dict[str, StageTrainer] = {}
    for name, is_stage2 in (("stage1", False), ("stage2", True)):
        model = models.get(name)
        if model is None:
            continue
-        spec = StageSpec.from_config(cfg, name, is_stage2, max(total_train_batches, 1))
-        if spec.generator == "wgan":
+        spec = StageSpec.from_config(cfg, name, is_stage2, max(total_train_batches, 1), sec_type_class_counts)
+        if build_objective(spec.generator).is_adversarial:
            critic = critics.get(name)
            assert critic is not None, (
                f"{name}_model.generator='wgan' requires a critic (see giant.model.network.build_critics)"
@@ -949,4 +1213,9 @@ def build_stage_trainers(
            trainers[name] = WGANStageTrainer(spec, model, critic, device)
        else:
            trainers[name] = FlowDDPMStageTrainer(spec, model, device)
+
+    stage2 = trainers.get("stage2")
+    stage1 = trainers.get("stage1")
+    if stage2 is not None and stage1 is not None and stage2.spec.stage1_context == "sampled":
+        stage2.attach_stage1(stage1)
    return trainers
@@ -82,7 +82,7 @@ def validate_marginals(
    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
@@ -103,7 +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 = [], []
@@ -115,11 +115,10 @@ 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)

@@ -130,10 +129,7 @@ def validate_marginals(
            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_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)))
@@ -141,6 +137,13 @@ 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
        )
+        # 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()

@@ -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.3.0"
+version = "0.4.0"
 description = "Geant4 step-function surrogate via conditional flow matching"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -26,7 +26,9 @@ dev = [
    "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",
@@ -47,16 +49,21 @@ 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", "scripts"]
+source = ["giant"]
 omit = ["*/legacy/*"]

 [tool.coverage.report]
@@ -70,7 +77,7 @@ requires = ["hatchling"]
 build-backend = "hatchling.build"

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

 [tool.uv]
 conflicts = [
@@ -101,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,75 +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 import config as gconfig
-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 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,
-    }
-    # Merged against DEFAULT_CONFIG (not a hand-rolled partial dict) so
-    # run_setup_stage always sees every key it might read (e.g.
-    # conditioning.particle.emb_dim, stage2_model.particle_type.target) at
-    # its real default, not silently missing/None — see issues.md Issue 1.
-    # 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 = gconfig.merge_cli_overrides(
-        gconfig.DEFAULT_CONFIG,
-        None,
-        {
-            "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.")
@@ -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,12 @@ from giant.analysis import reduce as R
 from giant.analysis.sources import (
    SYNTHETIC_TERMINATION_REASONS,
    Side,
+    open_side,
    physical_steps,
    secondaries,
+    secondaries_by_step,
 )
+from giant.data.loader import EVENT_ID_FILE_STRIDE


 def _rollout_frame() -> pl.LazyFrame:
@@ -102,6 +105,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 +160,19 @@ def test_secondaries_rollout_vs_reference_align():
    assert t["pdg"].to_list() == [22, 22]


+def test_secondaries_by_step_keys_each_secondary_to_its_emitting_step():
+    r = secondaries_by_step(_rollout_frame(), Side.rollout).collect()
+    assert r["pdg"].to_list() == [22]
+    # the rollout key is (event_id, parent_id, birth position) — the parent
+    # step's post_pos, copied verbatim onto the child's birth row.
+    assert r["step_key"][0] == {"event_id": 1, "parent_id": 0, "pre_x": 0.0, "pre_y": 0.0, "pre_z": 1.0}
+
+    t = secondaries_by_step(_reference_frame(), Side.reference).collect()
+    assert t["pdg"].to_list() == [22, 22]
+    # one row per emitting step; the empty-list step drops out entirely
+    assert [k["_row"] for k in t["step_key"]] == [0, 2]
+
+
 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]
@@ -0,0 +1,327 @@
+"""Tests for the rollout-YAML(s) → run-directory flow, compute, and merge."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pyarrow.parquet as pq
+import pytest
+import yaml
+
+from giant.analysis import (
+    RunMeta,
+    compute_one,
+    compute_reduced,
+    derive_run_dir,
+    load_rollout_yaml,
+    load_rollout_yamls,
+    merge_one,
+    prep,
+)
+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"
+    _write_rollout(rollout)
+    _reference_frame().collect().write_parquet(reference)
+
+    yaml_path = tmp_path / "run.yaml"
+    yaml_path.write_text(
+        yaml.safe_dump(
+            {
+                "prediction_id": "abcd1234ef",
+                "output": str(rollout),
+                "dataset": str(reference),
+                "checkpoint": "/ckpt/best.pt",
+                "kind": "rollout",
+                "energy_cutoff": 0.1,
+                "steps": 10,
+            }
+        )
+    )
+    return yaml_path
+
+
+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_yamls, run_dir: str | Path | None = None, chunks: int = 1, labels=None) -> Path:
+    """``prep`` with small test-sized context bins/sampling."""
+    return prep(
+        rollout_yamls,
+        run_dir,
+        n_chunks=chunks,
+        labels=labels,
+        n_energy_bins=2,
+        n_marginal_bins=8,
+        top_k_pdg=3,
+        sample_rows=1000,
+    )
+
+
+def test_load_rollout_yaml_requires_paths(tmp_path: Path):
+    bad = tmp_path / "bad.yaml"
+    bad.write_text(yaml.safe_dump({"output": "x.parquet"}))  # no dataset
+    with pytest.raises(ValueError):
+        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")
+
+
+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(
+        "/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")
+
+
+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])
+    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 [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)
+    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
+
+
+def test_reprep_clears_stale_partials_from_a_different_chunk_count(tmp_path: Path):
+    """Re-prepping with a different n_chunks must not leave old chunk
+    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)
+    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"
+    assert stale.exists()
+    (run_dir / "reduced").mkdir(exist_ok=True)
+    (run_dir / "reduced" / "marginal_edep.json").write_text("{}")
+
+    _prep([yaml_path], run_dir, chunks=1)
+
+    assert not stale.exists()
+    assert not (run_dir / "reduced" / "marginal_edep.json").exists()
+    assert (run_dir / "shared.json").exists()  # prep's own fresh output untouched
+
+
+def test_compute_one_from_run_dir(tmp_path: 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_timing_survives_plot_meta_to_compute_one(tmp_path: Path):
+    yaml_path = _write_inputs(tmp_path)
+    d = yaml.safe_load(yaml_path.read_text())
+    d["timing"] = {"us_per_step": 7.0, "write_us_per_step": 1.0}
+    yaml_path.write_text(yaml.safe_dump(d))
+
+    run_dir = _prep([yaml_path])
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    assert meta.rollouts[0]["plot_meta"]["timing"] == {"us_per_step": 7.0, "write_us_per_step": 1.0}
+
+    out = compute_one("eval_cost_per_step", run_dir)
+    reduced = Reduced(**Partial.load(out).data["reduced"])
+    assert reduced.kind == "bar"
+    assert reduced.payload["series"]["rollout"] == [7.0, 1.0, 8.0]
+
+
+def test_compute_reduced_explicit_paths(tmp_path: 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",
+        rollouts,
+        meta.reference,
+        run_dir / "shared.json",
+        tmp_path / "r.json",
+    )
+    assert Partial.load(out).id == "marginal_step_length"
+
+
+def test_merge_one_produces_reduced(tmp_path: 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["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)
+    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)
+
+
+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")])
+    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)
+    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))
+
+    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)
+    with pytest.raises(ValueError, match="out of range"):
+        compute_one("marginal_edep", run_dir, chunk_index=1)
+
+
+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
+
+    assert RUNTIME_SAFETY_MARGIN > 0
+    small = estimate_runtime_s("marginal_edep", 1_000)
+    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
@@ -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
@@ -6,14 +6,29 @@ 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,
+    _ks_statistic,
+)
 from giant.analysis.context import Context, build_context
+from giant.analysis.grouping import pdg_label
+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)
+    return build_context(
+        [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")
@@ -21,9 +36,20 @@ 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():
@@ -53,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"])


 # ---------------------------------------------------------------------------
@@ -98,7 +159,11 @@ 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),
+# sum-mergeable-with-a-zero-fill-denominator (sec_count_per_step{,_by_species}),
+# nested sum-merge into a scorecard (marginal_distance_summary), 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",
@@ -106,7 +171,11 @@ _CHUNK_EQUIVALENCE_IDS = [
    "shower_longitudinal",
    "leakage_fraction",
    "sec_count_per_species",
+    "sec_count_per_step",
+    "sec_count_per_step_by_species",
    "router_gating",
+    "marginal_distance_summary",
+    "shower_containment_depth_90",
 ]


@@ -128,21 +197,96 @@ 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 and 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_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_sec_count_per_step_counts_empty_steps(bundle):
+    spec = get_spec("sec_count_per_step")
+    r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
+    # reference: 3 steps, two of which emit exactly one secondary
+    assert r.payload["reference"][:2] == [1, 2]
+    # rollout: 4 physical steps, one of which emits a single secondary
+    assert r.payload["series"]["rollout"][:2] == [3, 1]
+    assert sum(r.payload["reference"]) == 3
+
+
+def test_sec_count_per_step_by_species_zero_row_is_per_species(bundle):
+    spec = get_spec("sec_count_per_step_by_species")
+    r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
+    cols = r.payload["col_labels"]
+    ref = r.payload["reference"]
+    g = cols.index(pdg_label(22))
+    # two reference steps emit one photon each; the third emits none
+    assert [row[g] for row in ref][:2] == [1, 2]
+    # every other species column is "no such secondary" on all 3 steps
+    for j, _ in enumerate(cols):
+        if j != g:
+            assert ref[0][j] == 3 and sum(row[j] for row in ref[1:]) == 0
+
+
+# ---------------------------------------------------------------------------
+# eval_cost_per_step
+# ---------------------------------------------------------------------------
+
+
+def test_eval_cost_per_step_unavailable_without_timing(bundle: Bundle):
+    # `bundle`'s RolloutSpec carries no `timing` -> no rollout to compare.
+    spec = get_spec("eval_cost_per_step")
+    r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
+    assert r.kind == "unavailable"
+    assert r.payload["note"]
+
+
+def test_eval_cost_per_step_bar_with_timing(ctx: Context):
+    spec = get_spec("eval_cost_per_step")
+    rs = RolloutSpec(
+        "rollout",
+        _rollout_frame(),
+        timing={"us_per_step": 12.5, "write_us_per_step": 2.5},
+    )
+    b = Bundle.open([rs], _reference_frame(), ctx)
+    r = spec.finalize([spec.compute_partial(b)], ctx)
+    assert r.kind == "bar"
+    assert r.payload["series"]["rollout"] == [12.5, 2.5, 15.0]
+    assert len(r.payload["reference"]) == 3
+    assert r.payload["log_y"] is True
+    assert "rollout" in r.meta["speedup_vs_geant4_total"]
@@ -14,6 +14,7 @@ from giant import config as gconfig
 from giant.checkpoint_io import (
    CheckpointCompatibilityError,
    InferenceContext,
+    apply_config_overrides,
    conditioning_axes,
    load_for_inference,
    stage_cfg,
@@ -73,6 +74,16 @@ def _write_checkpoint(tmp_path, model_cfg=None, ema: bool = False, **ckpt_overri
    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)
@@ -179,6 +190,40 @@ def test_onehot_particle_conditioning_with_topn_map_succeeds(tmp_path):
    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)

@@ -234,3 +279,122 @@ def test_stage_cfg_new_shape_returns_subdict():
 def test_stage_cfg_v02_flat_shape_returns_empty_dict():
    model_cfg = {"hidden_dim": 32, "n_blocks": 4}
    assert stage_cfg(model_cfg, "stage2") == {}
+
+
+# ---------------------------------------------------------------------------
+# config_overrides (gitea #87)
+# ---------------------------------------------------------------------------
+
+
+def _router_model_cfg() -> dict:
+    cfg = _model_cfg()
+    cfg["stage1_model"]["router"] = {"enabled": True, "type": "energy", "n_experts": 2}
+    return cfg
+
+
+def test_config_override_n_sec_sampling_changes_stage2_attribute(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(
+        checkpoint,
+        torch.device("cpu"),
+        "predict",
+        config_overrides={"stage2_model.n_sec.sampling": "sample"},
+    )
+    assert ctx.stage2 is not None
+    assert ctx.stage2.n_sec_sampling == "sample"
+    assert ctx.config_overrides == {"stage2_model.n_sec.sampling": "sample"}
+
+
+def test_config_override_ddpm_n_steps_changes_context_fields(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(
+        checkpoint,
+        torch.device("cpu"),
+        "predict",
+        config_overrides={"stage1_model.ddpm.n_steps": 42, "stage2_model.ddpm.n_steps": 7},
+    )
+    assert ctx.stage1_ddpm_steps == 42
+    assert ctx.stage2_ddpm_steps == 7
+
+
+def test_config_override_other_policy_changes_context_field(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(
+        checkpoint,
+        torch.device("cpu"),
+        "predict",
+        config_overrides={"stage2_model.particle_type.other_policy": "modal"},
+    )
+    assert ctx.other_policy == "modal"
+
+
+def test_config_override_router_temperature_changes_router_attribute(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, model_cfg=_router_model_cfg())
+    ctx = load_for_inference(
+        checkpoint,
+        torch.device("cpu"),
+        "predict",
+        config_overrides={"stage1_model.router.temperature": 1.5},
+    )
+    assert ctx.stage1 is not None
+    assert ctx.stage1.trunk.router.temperature == pytest.approx(1.5)
+
+
+def test_config_override_no_overrides_defaults_to_empty_dict(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+    assert ctx.config_overrides == {}
+
+
+def test_config_override_unknown_path_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    with pytest.raises(CheckpointCompatibilityError, match="not an inference-safe override"):
+        load_for_inference(
+            checkpoint,
+            torch.device("cpu"),
+            "predict",
+            config_overrides={"stage2_model.n_sec.typo": "sample"},
+        )
+
+
+def test_config_override_shape_bearing_key_raises_up_front(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    with pytest.raises(CheckpointCompatibilityError, match="not an inference-safe override"):
+        load_for_inference(
+            checkpoint,
+            torch.device("cpu"),
+            "predict",
+            config_overrides={"stage1_model.hidden_dim": 999},
+        )
+
+
+def test_config_override_bad_value_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    with pytest.raises(CheckpointCompatibilityError, match="must be one of"):
+        load_for_inference(
+            checkpoint,
+            torch.device("cpu"),
+            "predict",
+            config_overrides={"stage2_model.n_sec.sampling": "maybe"},
+        )
+
+
+def test_apply_config_overrides_no_overrides_returns_same_object():
+    cfg = _model_cfg()
+    assert apply_config_overrides(cfg, None) is cfg
+    assert apply_config_overrides(cfg, {}) is cfg
+
+
+def test_apply_config_overrides_migrates_legacy_flat_model_config_first():
+    legacy_cfg = {
+        "pdg_vocab": len(PDG_MAP),
+        "mat_vocab": len(MAT_MAP),
+        "hidden_dim": 32,
+        "n_blocks": 4,
+        "emb_dim": 8,
+        "dropout": 0.1,
+        "k_max": 5,
+    }
+    merged = apply_config_overrides(legacy_cfg, {"stage1_model.ddpm.n_steps": 10})
+    assert merged["stage1_model"]["ddpm"]["n_steps"] == 10
+    assert merged["stage1_model"]["hidden_dim"] == 32
@@ -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()
@@ -172,3 +172,43 @@ def test_predict_exits_1_on_checkpoint_missing_model_config(tmp_path):

    assert result.exit_code == 1
    assert "checkpoint has no model_config" in result.output
+
+
+# ---------------------------------------------------------------------------
+# --set (gitea #87)
+# ---------------------------------------------------------------------------
+
+
+def test_predict_set_flag_without_equals_exits_1(tmp_path):
+    checkpoint = tmp_path / "missing.pt"
+
+    result = runner.invoke(
+        app,
+        ["predict", "dummy.parquet", "--checkpoint", str(checkpoint), "--set", "sampling"],
+    )
+
+    assert result.exit_code == 1
+    assert "must be 'dotted.path=value'" in result.output
+
+
+def test_predict_set_flag_disallowed_path_surfaces_compat_error(tmp_path):
+    checkpoint = tmp_path / "ckpt.pt"
+    torch.save(
+        {"model_config": {"stage1_model": {}, "stage2_model": {}}, "sec_decoder": {}, "normalizer": {"sec_phys": {}}},
+        checkpoint,
+    )
+
+    result = runner.invoke(
+        app,
+        [
+            "predict",
+            "dummy.parquet",
+            "--checkpoint",
+            str(checkpoint),
+            "--set",
+            "stage1_model.hidden_dim=999",
+        ],
+    )
+
+    assert result.exit_code == 1
+    assert "not an inference-safe override" in result.output
@@ -8,11 +8,52 @@ from __future__ import annotations
 import torch
 from typer.testing import CliRunner

-from giant.cli import app
+from giant.cli import _build_rollout_timing, app

 runner = CliRunner()


+def test_build_rollout_timing_excludes_synthetic_rows_from_per_step_cost():
+    # 100 rows total, 30 of them synthetic termination markers (escape) ->
+    # us_per_step should be normalized over the 70 physical rows only, the
+    # same unit giant.analysis.geant4_reference measures Geant4 in.
+    timing = _build_rollout_timing(
+        setup_s=1.0,
+        rollout_s=10.0,
+        write_s=2.0,
+        n_rows=100,
+        termination_reason_counts={"escaped": 30, "natural_end": 70},
+        n_seed_events=5,
+        device="cpu",
+        torch_threads=4,
+    )
+    assert timing["n_rows"] == 100
+    assert timing["n_physical_rows"] == 70
+    assert timing["n_physical_rows"] < timing["n_rows"]
+    assert timing["sample_s"] == 8.0  # rollout_s - write_s
+    assert timing["us_per_step"] == 8.0 / 70 * 1e6
+    assert timing["write_us_per_step"] == 2.0 / 70 * 1e6
+    assert timing["ms_per_event"] == 10.0 / 5 * 1e3
+    assert timing["device"] == "cpu" and timing["torch_threads"] == 4
+
+
+def test_build_rollout_timing_handles_zero_physical_rows_and_events():
+    timing = _build_rollout_timing(
+        setup_s=1.0,
+        rollout_s=1.0,
+        write_s=0.0,
+        n_rows=5,
+        termination_reason_counts={"escaped": 5},
+        n_seed_events=0,
+        device="cpu",
+        torch_threads=1,
+    )
+    assert timing["n_physical_rows"] == 0
+    assert timing["us_per_step"] is None
+    assert timing["write_us_per_step"] is None
+    assert timing["ms_per_event"] is None
+
+
 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)
@@ -31,3 +72,28 @@ def test_rollout_exits_1_on_checkpoint_missing_model_config(tmp_path):

    assert result.exit_code == 1
    assert "checkpoint has no model_config" in result.output
+
+
+def test_rollout_set_flag_disallowed_path_surfaces_compat_error(tmp_path):
+    checkpoint = tmp_path / "ckpt.pt"
+    torch.save(
+        {"model_config": {"stage1_model": {}, "stage2_model": {}}, "sec_decoder": {}, "normalizer": {"sec_phys": {}}},
+        checkpoint,
+    )
+
+    result = runner.invoke(
+        app,
+        [
+            "rollout",
+            "dummy.parquet",
+            "--checkpoint",
+            str(checkpoint),
+            "--geometry",
+            "dummy_geometry.pkl",
+            "--set",
+            "stage2_model.n_sec.typo=sample",
+        ],
+    )
+
+    assert result.exit_code == 1
+    assert "not an inference-safe override" in result.output
@@ -20,7 +20,7 @@ def _invoke_and_capture_cfg(monkeypatch, tmp_path: Path, args: list[str]) -> dic
    def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
        captured["cfg"] = cfg

-    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr("giant.pipeline.run_train_job", _fake_run_train_job)

    result = runner.invoke(
        cli.app,
@@ -41,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,
@@ -53,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

@@ -96,8 +97,35 @@ def test_wgan_knobs_split_per_stage(monkeypatch, tmp_path):
    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)
+    monkeypatch.setattr("giant.pipeline.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"],
@@ -112,7 +140,7 @@ def test_out_dir_resolution_prefers_explicit_out_over_resume(monkeypatch, tmp_pa
    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.setattr("giant.pipeline.run_train_job", _fake_run_train_job)

    resume_dir = tmp_path / "resumed_run"
    resume_dir.mkdir()
@@ -133,7 +161,7 @@ def test_out_dir_resolution_falls_back_to_resume_parent(monkeypatch, tmp_path):
    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.setattr("giant.pipeline.run_train_job", _fake_run_train_job)

    resume_dir = tmp_path / "resumed_run"
    resume_dir.mkdir()
@@ -150,7 +178,7 @@ def test_out_dir_resolution_defaults_when_neither_out_nor_resume_given(monkeypat
    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.setattr("giant.pipeline.run_train_job", _fake_run_train_job)
    monkeypatch.chdir(tmp_path)

    result = runner.invoke(cli.app, ["train", "dummy.parquet"])
@@ -164,7 +192,7 @@ def test_batch_size_auto_estimates_and_echoes(monkeypatch, tmp_path):
    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("giant.pipeline.run_train_job", _fake_run_train_job)
    monkeypatch.setattr(cli.gconfig, "estimate_batch_size", lambda hidden_dim, n_blocks, device: 123)

    result = runner.invoke(
@@ -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")
@@ -1,293 +0,0 @@
-"""Tests for the rollout-YAML → run-directory flow, compute, and submit."""
-
-from __future__ import annotations
-
-import sys
-from pathlib import Path
-
-import pyarrow.parquet as pq
-import pytest
-import yaml
-
-from giant.analysis import (
-    RunMeta,
-    SubmitConfig,
-    catalog_ids,
-    compute_one,
-    compute_reduced,
-    derive_run_dir,
-    load_rollout_yaml,
-    merge_one,
-    prep,
-    write_submit,
-)
-from giant.analysis.catalog import get_spec
-from giant.analysis.condor 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_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)
-    _reference_frame().collect().write_parquet(reference)
-
-    yaml_path = tmp_path / "run.yaml"
-    yaml_path.write_text(
-        yaml.safe_dump(
-            {
-                "prediction_id": "abcd1234ef",
-                "output": str(rollout),
-                "dataset": str(reference),
-                "checkpoint": "/ckpt/best.pt",
-                "kind": "rollout",
-                "energy_cutoff": 0.1,
-                "steps": 10,
-            }
-        )
-    )
-    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 _prep(rollout_yaml: Path, run_dir: str | Path | None = None, chunks: int = 1) -> Path:
-    """``prep`` with small test-sized context bins/sampling."""
-    return prep(
-        rollout_yaml,
-        run_dir,
-        n_chunks=chunks,
-        n_energy_bins=2,
-        n_marginal_bins=8,
-        top_k_pdg=3,
-        sample_rows=1000,
-    )
-
-
-def test_load_rollout_yaml_requires_paths(tmp_path: Path):
-    bad = tmp_path / "bad.yaml"
-    bad.write_text(yaml.safe_dump({"output": "x.parquet"}))  # no dataset
-    with pytest.raises(ValueError):
-        load_rollout_yaml(bad)
-
-
-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")
-
-
-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(
-        "/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")
-
-
-def test_prep_lays_out_run_dir(tmp_path: Path):
-    yaml_path = _write_inputs(tmp_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 "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_splits_rows_per_chunk(tmp_path: Path):
-    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
-
-
-def test_reprep_clears_stale_partials_from_a_different_chunk_count(tmp_path: Path):
-    """Re-prepping with a different n_chunks must not leave old chunk
-    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)
-    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"
-    assert stale.exists()
-    (run_dir / "reduced").mkdir(exist_ok=True)
-    (run_dir / "reduced" / "marginal_edep.json").write_text("{}")
-
-    _prep(yaml_path, run_dir, chunks=1)
-
-    assert not stale.exists()
-    assert not (run_dir / "reduced" / "marginal_edep.json").exists()
-    assert (run_dir / "shared.json").exists()  # prep's own fresh output untouched
-
-
-def test_compute_one_from_run_dir(tmp_path: 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
-
-
-def test_compute_reduced_explicit_paths(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path))
-    meta = RunMeta.load(run_dir / "run_meta.json")
-    out = compute_reduced(
-        "marginal_step_length",
-        meta.rollout,
-        meta.reference,
-        run_dir / "shared.json",
-        tmp_path / "r.json",
-    )
-    assert Partial.load(out).id == "marginal_step_length"
-
-
-def test_merge_one_produces_reduced(tmp_path: 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
-
-
-def test_merge_one_fails_loudly_on_missing_chunk(tmp_path: Path):
-    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)
-
-
-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"))
-    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)
-    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))
-
-    assert chunked.payload == unchunked.payload
-
-
-def test_compute_reduced_rejects_out_of_range_chunk(tmp_path: Path):
-    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
-
-    assert RUNTIME_SAFETY_MARGIN > 0
-    small = estimate_runtime_s("marginal_edep", 1_000)
-    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
@@ -51,9 +51,13 @@ def test_giant_config_to_dict_matches_default_config():
        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,
@@ -75,6 +79,70 @@ def test_stage2_model_config_defaults_match_documented_v030_intent():
    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)
@@ -95,10 +163,48 @@ def test_stage1_router_config_has_no_tie_to_stage1_key():
    assert "tie_to_stage1" not in gconfig.RouterConfig().to_dict()


-def test_n_sec_config_extra_round_trips_legacy_owner():
-    n_sec = gconfig.NSecConfig.from_dict({"mode": "head", "legacy_owner": "stage1"})
-    assert n_sec.legacy_owner == "stage1"
-    assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "legacy_owner": "stage1"}
+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", "sampling": "greedy"}
+
+
+def test_n_sec_config_sampling_defaults_to_greedy():
+    assert gconfig.NSecConfig().sampling == "greedy"
+
+
+def test_n_sec_config_sampling_round_trips():
+    n_sec = gconfig.NSecConfig.from_dict({"mode": "stop_token", "sampling": "sample"})
+    assert n_sec.sampling == "sample"
+    assert n_sec.to_dict()["sampling"] == "sample"
+
+
+def test_n_sec_config_stop_sampling_alias_still_honored():
+    """gitea #86: stop_sampling was renamed to sampling; old checkpoints'
+    model_config still carries the old key and must keep working."""
+    n_sec = gconfig.NSecConfig.from_dict({"mode": "stop_token", "stop_sampling": "sample"})
+    assert n_sec.sampling == "sample"
+    assert "stop_sampling" not in n_sec.to_dict()
+
+
+def test_n_sec_config_sampling_key_wins_over_stop_sampling_alias():
+    n_sec = gconfig.NSecConfig.from_dict({"sampling": "sample", "stop_sampling": "greedy"})
+    assert n_sec.sampling == "sample"
+
+
+def test_migrate_config_renames_stop_sampling_key():
+    cfg = {
+        "meta": {"config_version": gconfig.CONFIG_VERSION},
+        "stage2_model": {"n_sec": {"stop_sampling": "sample"}},
+    }
+    migrated = gconfig.migrate_config(cfg)
+    assert gconfig._get_path(migrated, "stage2_model.n_sec.sampling") == "sample"
+    assert gconfig._get_path(migrated, "stage2_model.n_sec.stop_sampling") is None


 # ---------------------------------------------------------------------------
@@ -629,6 +735,42 @@ 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():
    """The particle and material conditioning axes are configured
    independently and may mix freely — e.g. material
@@ -673,13 +815,166 @@ 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_n_sec_sampling_rejected():
+    cfg = _cfg_with(**{"stage2_model.n_sec.sampling": "bogus"})
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "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():
@@ -733,6 +1028,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(
@@ -786,11 +1106,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",
        }
@@ -861,6 +1182,41 @@ def test_validate_config_keys_skips_meta_section():
    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")
@@ -907,6 +1263,11 @@ def test_overrides_from_flags_train_block_passthrough():
    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"),
    [
@@ -1008,6 +1369,38 @@ def test_overrides_from_flags_wgan_knobs_split_per_stage(shared, stage1_specific
    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)
 # ---------------------------------------------------------------------------
@@ -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
@@ -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
@@ -124,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",
@@ -167,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():
@@ -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)
@@ -195,6 +195,10 @@ def test_build_topn_map_from_files_keeps_most_frequent(tmp_path):
    assert m.class_map["G4_Fe"] == 2  # "other" (n_classes - 1)
    assert m.class_map["G4_Pb"] == 2
    assert m.other_members == {"G4_Fe": 2, "G4_Pb": 1}
+    # class_counts (gitea #44): per resulting index, "other" is the sum of
+    # everything folded into it (2 + 1 = 3), and the total equals row count.
+    assert m.class_counts == {0: 5, 1: 3, 2: 3}
+    assert sum(m.class_counts.values()) == len(materials)


 def test_build_topn_map_from_files_fewer_values_than_n_classes(tmp_path):
@@ -205,6 +209,8 @@ def test_build_topn_map_from_files_fewer_values_than_n_classes(tmp_path):

    assert m.class_map == {"G4_AIR": 0, "PbWO4": 1}
    assert m.other_members == {}
+    # No "other" bucket ever populated -> no entry for its index either.
+    assert m.class_counts == {0: 1, 1: 1}


 def test_build_pdg_topn_map_from_files_pools_primary_and_secondary_pdg(tmp_path):
@@ -224,6 +230,7 @@ def test_build_pdg_topn_map_from_files_pools_primary_and_secondary_pdg(tmp_path)
    # pooled: 11 -> 5, 22 -> 1 (primary) + 10 (secondary) = 11
    assert m.class_map[22] == 0
    assert m.class_map[11] == 1
+    assert m.class_counts == {0: 11, 1: 5}


 def test_build_pdg_topn_map_from_files_missing_sec_pdg_list_column(tmp_path):
@@ -130,7 +130,7 @@ def _run_migration_check(mode: str, conditioning: str) -> None:
    new_stage1, new_stage2 = new_models["stage1"], new_models["stage2"]
    assert isinstance(new_stage1, net.Stage1Model)
    assert isinstance(new_stage2, net.Stage2OneShot)
-    # legacy_owner="stage1": n_sec lives on stage1, not stage2, for a
+    # n_sec.owner="stage1": n_sec lives on stage1, not stage2, for a
    # migrated v0.2 checkpoint.
    assert new_stage1.n_sec_head is not None
    assert new_stage2.n_sec_head is None
@@ -177,7 +177,7 @@ def test_migration_wgan_physical():

 def test_migrate_legacy_model_config_shape():
    """_migrate_legacy_model_config produces the nested shape build_models
-    expects, with the legacy_owner marker set so build_models routes the
+    expects, with the n_sec.owner marker set so build_models routes the
    n_sec head back onto stage 1."""
    legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
    migrated = net._migrate_legacy_model_config(legacy_cfg)
@@ -187,7 +187,7 @@ def test_migrate_legacy_model_config_shape():
    assert migrated["conditioning"]["particle"]["n_layers"] == 2
    assert migrated["conditioning"]["material"]["n_layers"] == 2
    assert migrated["stage1_model"]["hidden_dim"] == HIDDEN_DIM
-    assert migrated["stage2_model"]["n_sec"]["legacy_owner"] == "stage1"
+    assert migrated["stage2_model"]["n_sec"]["owner"] == "stage1"
    assert migrated["stage2_model"]["decoder"] == "one_shot"


@@ -278,6 +278,6 @@ def test_build_models_accepts_new_nested_shape_unchanged():
    models = net.build_models(cfg)
    assert isinstance(models["stage1"], net.Stage1Model)
    assert isinstance(models["stage2"], net.Stage2OneShot)
-    # Fresh v0.3.0 config, no legacy_owner: n_sec lives on stage 2.
+    # Fresh v0.3.0 config, n_sec.owner defaults to "stage2": n_sec lives on stage 2.
    assert models["stage1"].n_sec_head is None
    assert models["stage2"].n_sec_head is not None
@@ -0,0 +1,151 @@
+"""Tests for `giant model summary` (gitea #46)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+from typer.testing import CliRunner
+
+from giant import config as gconfig
+from giant.cli import app
+from giant.materials import MATERIAL_PROPERTIES
+from giant.model.summary import _NOT_BUILD_TIME, _built_modules, _vocab_caveats, summarize_model
+
+INFERENCE_OVERRIDES = gconfig.INFERENCE_OVERRIDES
+
+runner = CliRunner()
+
+_PDG_VOCAB = 300
+_MAT_VOCAB = len(MATERIAL_PROPERTIES)
+
+
+def _cfg(overrides: dict | None = None) -> dict:
+    return gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, overrides or {})
+
+
+@pytest.fixture(scope="module")
+def default_summary():
+    return summarize_model(_cfg(), pdg_vocab=_PDG_VOCAB, mat_vocab=_MAT_VOCAB)
+
+
+def test_default_config_builds_both_stages_with_a_real_tree(default_summary):
+    assert set(default_summary.modules) >= {"stage1", "stage2"}
+    for module in default_summary.modules.values():
+        assert sum(p.numel() for p in module.parameters()) > 0
+    stage1 = default_summary.modules["stage1"]
+    assert hasattr(stage1, "cond_enc")
+    assert hasattr(stage1, "trunk")
+    assert {"input_proj", "blocks", "out_proj"} <= {n for n, _ in stage1.trunk.named_children()}
+
+
+def test_every_in_scope_leaf_is_classified(default_summary):
+    in_scope = {
+        p
+        for p in gconfig.leaf_paths(gconfig.DEFAULT_CONFIG)
+        if p.split(".", 1)[0] in ("conditioning", "stage1_model", "stage2_model")
+    }
+    classified = set(default_summary.consumed) | set(default_summary.inert) | set(default_summary.elsewhere)
+    assert classified == in_scope
+
+
+def test_not_build_time_allow_list_has_no_stale_entries():
+    in_scope = set(gconfig.leaf_paths(gconfig.DEFAULT_CONFIG))
+    stale = set(_NOT_BUILD_TIME) - in_scope
+    assert not stale, f"_NOT_BUILD_TIME entries no longer in DEFAULT_CONFIG: {sorted(stale)}"
+
+
+def test_inference_overrides_allow_list_has_no_stale_entries():
+    in_scope = set(gconfig.leaf_paths(gconfig.DEFAULT_CONFIG))
+    stale = set(INFERENCE_OVERRIDES) - in_scope
+    assert not stale, f"INFERENCE_OVERRIDES entries no longer in DEFAULT_CONFIG: {sorted(stale)}"
+
+
+def test_default_config_overridable_lists_every_allowlisted_path(default_summary):
+    assert set(default_summary.overridable) == set(INFERENCE_OVERRIDES)
+
+
+def test_router_disabled_by_default_so_its_fields_are_inert(default_summary):
+    assert "stage1_model.router.n_experts" in default_summary.inert
+    assert "stage1_model.router.temperature" in default_summary.inert
+
+
+def test_markov_history_leaves_attention_dims_inert_but_history_itself_consumed(default_summary):
+    assert "stage2_model.autoregressive.attn_n_heads" in default_summary.inert
+    assert "stage2_model.autoregressive.attn_n_layers" in default_summary.inert
+    assert "stage2_model.autoregressive.history" in default_summary.consumed
+
+
+def test_single_literal_branch_fields_are_correctly_seen_as_consumed(default_summary):
+    """Regression guard: n_sec.owner ("stage2"), n_sec.mode ("head") and
+    particle_type.target ("onehot") each branch as `== "one specific other
+    literal"` in giant/model/builders.py|models.py. A naive single generic
+    sentinel probe lands in the same "not that literal" bucket as the
+    current value and never crosses the boundary that actually matters --
+    this is exactly what _STRING_ALTERNATIVES exists to fix."""
+    assert "stage2_model.n_sec.owner" in default_summary.consumed
+    assert "stage2_model.n_sec.mode" in default_summary.consumed
+    assert "stage2_model.particle_type.target" in default_summary.consumed
+
+
+def test_stage1_wgan_generator_swaps_flow_time_dim_for_critic_dims():
+    summary = summarize_model(_cfg({"stage1_model": {"generator": "wgan"}}), pdg_vocab=_PDG_VOCAB, mat_vocab=_MAT_VOCAB)
+    assert "stage1_model.flow.time_dim" in summary.inert
+    assert "stage1_model.wgan.noise_dim" in summary.consumed
+    assert "stage1_model.wgan.critic_hidden_dim" in summary.consumed
+
+
+def test_stage2_one_shot_decoder_makes_autoregressive_block_inert():
+    summary = summarize_model(
+        _cfg({"stage2_model": {"decoder": "one_shot"}}), pdg_vocab=_PDG_VOCAB, mat_vocab=_MAT_VOCAB
+    )
+    assert "stage2_model.autoregressive.history" in summary.inert
+    assert "history_encoder" not in {n for n, _ in summary.modules["stage2"].named_children()}
+
+
+def test_energy_router_enabled_consumes_core_fields_but_not_process_only_fields():
+    summary = summarize_model(
+        _cfg({"stage1_model": {"router": {"enabled": True, "type": "energy", "n_experts": 4}}}),
+        pdg_vocab=_PDG_VOCAB,
+        mat_vocab=_MAT_VOCAB,
+    )
+    assert "stage1_model.router.n_experts" in summary.consumed
+    assert "stage1_model.router.temperature" in summary.consumed
+    # emb_dim/hidden_dim are pdg/process-router-only kwargs -- build_router's
+    # signature filter drops them for an energy router.
+    assert "stage1_model.router.hidden_dim" in summary.inert
+    assert "stage1_model.router.emb_dim" in summary.inert
+
+
+def test_vocab_caveat_text_for_embedding_particle_conditioning():
+    cfg = _cfg({"conditioning": {"particle": {"type": "embedding"}}})
+    caveats = _vocab_caveats(cfg)
+    assert any("pdg_vocab" in c and "embedding" in c for c in caveats)
+    assert not any("mat_vocab" in c for c in caveats)
+
+
+def test_pdg_vocab_flag_changes_embedding_table_size():
+    cfg = _cfg({"conditioning": {"particle": {"type": "embedding"}}})
+    small = _built_modules(cfg, pdg_vocab=10, mat_vocab=_MAT_VOCAB)
+    big = _built_modules(cfg, pdg_vocab=1000, mat_vocab=_MAT_VOCAB)
+    assert big["stage1"].cond_enc.pdg_emb.weight.numel() > small["stage1"].cond_enc.pdg_emb.weight.numel()
+
+
+def test_invalid_combo_exits_nonzero_with_validate_config_message(tmp_path: Path):
+    config_path = tmp_path / "bad.toml"
+    config_path.write_text('[meta]\nconfig_version = 3\n\n[stage2_model.particle_type]\ntarget = "embedding"\n')
+    result = runner.invoke(app, ["model", "summary", "--config", str(config_path)])
+    assert result.exit_code == 1
+    assert "requires conditioning.particle.type = 'embedding'" in result.output
+
+
+def test_cli_default_smoke():
+    result = runner.invoke(app, ["model", "summary"])
+    assert result.exit_code == 0, result.output
+    assert "stage1" in result.output
+    assert "stage2" in result.output
+    assert "parameters" in result.output
+    assert "trunk" in result.output
+    assert "inert under this config" in result.output
+    assert "inference-overridable without retraining" in result.output
+    assert "stage2_model.n_sec.sampling" in result.output
@@ -5,24 +5,32 @@ import torch
 from giant import config as gconfig
 from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
 from giant.model.network import (
+    HISTORY_REGISTRY,
    AttentionHistory,
    ConditionEncoder,
+    CriticModel,
+    FilmResBlock,
+    HistoryEncoder,
+    LinearTrunk,
    MarkovHistory,
+    NoHistory,
    SinusoidalEmbedding,
    Stage1Model,
    Stage2Autoregressive,
    Stage2OneShot,
+    StageModel,
    build_critics,
+    build_history,
    build_models,
-    cat_col_layout,
+    build_objective,
    stage2_trunk_sec_dim,
    stage2_type_dim,
 )

-PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
-MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
-ONEHOT_PARTICLE_CFG = {"type": "onehot", "emb_dim": 6, "n_layers": 1}
-ONEHOT_MATERIAL_CFG = {"type": "onehot", "emb_dim": 4, "n_layers": 1}
+PARTICLE_CFG = gconfig.ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
+MATERIAL_CFG = gconfig.ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
+ONEHOT_PARTICLE_CFG = gconfig.ConditioningAxisConfig(type="onehot", emb_dim=6, n_layers=1)
+ONEHOT_MATERIAL_CFG = gconfig.ConditioningAxisConfig(type="onehot", emb_dim=4, n_layers=1)


 def test_sinusoidal_embedding_shape():
@@ -90,48 +98,72 @@ def test_stage1_model_no_n_sec_head_by_default():
    assert model.n_sec_head is None


-# --- cat_col_layout / stage2_type_dim / stage2_trunk_sec_dim ---------------
+def test_stage1_model_n_sec_head_default_cfg_matches_pre_gitea_36_shape():
+    """No n_sec_head_cfg given must reproduce the old hardcoded
+    hidden_dim // 2, one-hidden-layer architecture exactly (gitea #36)."""
+    model = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=40,
+        cond_out_dim=12,
+        n_sec_head_k_max=15,
+    )
+    assert model.n_sec_head is not None
+    assert len(model.n_sec_head) == 3
+    assert model.n_sec_head[0].in_features == 12
+    assert model.n_sec_head[0].out_features == 20  # hidden_dim // 2
+    assert model.n_sec_head[2].out_features == 16  # k_max + 1


-def test_cat_col_layout_neither_onehot():
-    assert cat_col_layout("physical", "embedding") == (None, None)
+def test_stage1_model_n_sec_head_cfg_controls_hidden_width_and_depth():
+    model = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=32,
+        cond_out_dim=16,
+        n_sec_head_k_max=15,
+        n_sec_head_cfg={"hidden_ratio": 0.25, "depth": 1},
+    )
+    assert model.n_sec_head is not None
+    assert len(model.n_sec_head) == 1
+    assert model.n_sec_head[0].in_features == 16
+    assert model.n_sec_head[0].out_features == 16


-def test_cat_col_layout_particle_only():
-    assert cat_col_layout("onehot", "physical") == (2, None)
-
-
-def test_cat_col_layout_material_only():
-    assert cat_col_layout("physical", "onehot") == (None, 2)
-
-
-def test_cat_col_layout_both_onehot_particle_then_material():
-    assert cat_col_layout("onehot", "onehot") == (2, 3)
+# --- stage2_type_dim / stage2_trunk_sec_dim --------------------------------
+# (the cond_cat column-layout tests live in tests/test_cond_layout.py)


 def test_stage2_type_dim_physical_is_particle_phys_dim():
-    assert stage2_type_dim({"target": "physical"}, emb_dim=16) == PARTICLE_PHYS_DIM
+    assert stage2_type_dim(gconfig.ParticleTypeConfig(target="physical"), emb_dim=16) == PARTICLE_PHYS_DIM


 def test_stage2_type_dim_onehot_and_embedding_are_emb_dim():
-    assert stage2_type_dim({"target": "onehot"}, emb_dim=16) == 16
-    assert stage2_type_dim({"target": "embedding"}, emb_dim=16) == 16
+    assert stage2_type_dim(gconfig.ParticleTypeConfig(target="onehot"), emb_dim=16) == 16
+    assert stage2_type_dim(gconfig.ParticleTypeConfig(target="embedding"), emb_dim=16) == 16


 def test_stage2_trunk_sec_dim_physical_matches_v02_sec_dim():
    k_max = 15
-    assert stage2_trunk_sec_dim({"target": "physical"}, "flow", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM
-    assert stage2_trunk_sec_dim({"target": "physical"}, "wgan", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM
+    physical = gconfig.ParticleTypeConfig(target="physical")
+    assert stage2_trunk_sec_dim(physical, "flow", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM
+    assert stage2_trunk_sec_dim(physical, "wgan", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM


 def test_stage2_trunk_sec_dim_onehot_wgan_folds_type_in():
    k_max = 15
-    assert stage2_trunk_sec_dim({"target": "onehot"}, "wgan", k_max, emb_dim=16) == k_max * (CONT_SLOT_DIM + 16)
+    onehot = gconfig.ParticleTypeConfig(target="onehot")
+    assert stage2_trunk_sec_dim(onehot, "wgan", k_max, emb_dim=16) == k_max * (CONT_SLOT_DIM + 16)


 def test_stage2_trunk_sec_dim_onehot_flow_excludes_type():
    k_max = 15
-    assert stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim=16) == k_max * CONT_SLOT_DIM
+    onehot = gconfig.ParticleTypeConfig(target="onehot")
+    assert stage2_trunk_sec_dim(onehot, "flow", k_max, emb_dim=16) == k_max * CONT_SLOT_DIM


 # --- ConditionEncoder onehot mode -------------------------------------------
@@ -139,8 +171,8 @@ def test_stage2_trunk_sec_dim_onehot_flow_excludes_type():

 def test_condition_encoder_onehot_forward_shape_and_gradients():
    B = 8
-    particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"])
-    material_emb_dim = int(ONEHOT_MATERIAL_CFG["emb_dim"])
+    particle_emb_dim = ONEHOT_PARTICLE_CFG.emb_dim
+    material_emb_dim = ONEHOT_MATERIAL_CFG.emb_dim
    enc = ConditionEncoder(
        pdg_vocab=5,
        mat_vocab=3,
@@ -171,12 +203,12 @@ def test_condition_encoder_onehot_is_a_true_one_hot_vector():
    verify the concatenated input segment really is one-hot, not e.g. an
    accidentally-learned embedding."""
    B = 4
-    particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"])
+    particle_emb_dim = ONEHOT_PARTICLE_CFG.emb_dim
    enc = ConditionEncoder(
        pdg_vocab=5,
        mat_vocab=3,
        particle_cfg=ONEHOT_PARTICLE_CFG,
-        material_cfg={"type": "physical", "emb_dim": 4, "n_layers": 1},
+        material_cfg=gconfig.ConditioningAxisConfig(type="physical", emb_dim=4, n_layers=1),
        out_dim=16,
    )
    cond_cont = torch.zeros(B, COND_DIM)
@@ -198,13 +230,11 @@ def test_condition_encoder_onehot_is_a_true_one_hot_vector():


 def _build_stage2(target: str, generator: str, emb_dim: int = 6) -> Stage2OneShot:
-    particle_cfg = {"type": "physical", "emb_dim": emb_dim, "n_layers": 1}
-    if target != "physical":
-        particle_cfg = dict(particle_cfg)
-        if target == "embedding":
-            particle_cfg["type"] = "embedding"
+    particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=emb_dim, n_layers=1)
+    if target == "embedding":
+        particle_cfg = gconfig.ConditioningAxisConfig(type="embedding", emb_dim=emb_dim, n_layers=1)
    k_max = 5
-    sec_dim = stage2_trunk_sec_dim({"target": target}, generator, k_max, emb_dim)
+    sec_dim = stage2_trunk_sec_dim(gconfig.ParticleTypeConfig(target=target), generator, k_max, emb_dim)
    return Stage2OneShot(
        pdg_vocab=5,
        mat_vocab=3,
@@ -217,7 +247,7 @@ def _build_stage2(target: str, generator: str, emb_dim: int = 6) -> Stage2OneSho
        sec_dim=sec_dim,
        generator=generator,
        k_max=k_max,
-        particle_type_cfg={"target": target, "lambda": 1.0},
+        particle_type_cfg=gconfig.ParticleTypeConfig(target=target),
    )


@@ -265,6 +295,38 @@ def test_stage2_oneshot_predict_type_raises_when_no_type_head():
        pass


+def test_stage2_oneshot_n_sec_head_and_type_head_cfg_control_hidden_width_and_depth():
+    """gitea #36: n_sec_head_cfg/type_head_cfg are independently tunable."""
+    k_max, emb_dim = 5, 6
+    particle_cfg = gconfig.ConditioningAxisConfig(type="onehot", emb_dim=emb_dim, n_layers=1)
+    sec_dim = stage2_trunk_sec_dim(gconfig.ParticleTypeConfig(target="onehot"), "flow", k_max, emb_dim)
+    model = Stage2OneShot(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=particle_cfg,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=40,
+        n_res_blocks=1,
+        cond_out_dim=12,
+        context_dim=8,
+        sec_dim=sec_dim,
+        generator="flow",
+        k_max=k_max,
+        particle_type_cfg=gconfig.ParticleTypeConfig(target="onehot"),
+        n_sec_head_cfg={"hidden_ratio": 0.25, "depth": 1},
+        type_head_cfg={"hidden_ratio": 0.75, "depth": 2},
+    )
+    assert model.n_sec_head is not None
+    assert len(model.n_sec_head) == 1
+    assert model.n_sec_head[0].in_features == 12
+    assert model.n_sec_head[0].out_features == k_max + 1
+
+    assert model.type_head is not None
+    assert len(model.type_head) == 3
+    assert model.type_head[0].out_features == 30  # round(40 * 0.75)
+    assert model.type_head[2].out_features == k_max * emb_dim
+
+
 def test_stage2_oneshot_forward_shape_onehot_wgan():
    B, k_max, emb_dim = 4, 5, 6
    model = _build_stage2("onehot", "wgan", emb_dim=emb_dim)
@@ -288,6 +350,33 @@ def test_stage2_oneshot_forward_shape_onehot_flow_excludes_type():
    assert out.shape == (B, k_max * CONT_SLOT_DIM)


+def test_stage2_oneshot_particle_type_n_classes_overrides_conditioning_emb_dim():
+    """gitea #29: stage2_model.particle_type.n_classes, not
+    conditioning.particle.emb_dim, sizes the onehot type_head/type_dim when
+    explicitly set — the two used to be silently the same number."""
+    k_max = 5
+    particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=6, n_layers=1)
+    particle_type_cfg = gconfig.ParticleTypeConfig(target="onehot", n_classes=20)
+    sec_dim = stage2_trunk_sec_dim(particle_type_cfg, "flow", k_max, 20)
+    model = Stage2OneShot(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=particle_cfg,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        cond_out_dim=16,
+        context_dim=8,
+        sec_dim=sec_dim,
+        generator="flow",
+        k_max=k_max,
+        particle_type_cfg=particle_type_cfg,
+    )
+    assert model.type_dim == 20  # not particle_cfg.emb_dim == 6
+    assert model.type_head is not None
+    assert model.type_head[-1].out_features == k_max * 20
+
+
 # --- MarkovHistory -----------------------------------------------------------


@@ -380,6 +469,84 @@ def test_attention_history_step_matches_forward():
    assert torch.allclose(stepped, expected, atol=1e-5)


+# --- NoHistory (gitea #45) ----------------------------------------------------
+
+
+def test_no_history_shape():
+    hist = NoHistory(in_dim=7, out_dim=12)
+    B, K = 3, 5
+    feat = torch.randn(B, K, 7)
+    has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
+    out = hist(feat, has_prev)
+    assert out.shape == (B, K, 12)
+
+
+def test_no_history_ignores_feat_and_has_prev():
+    hist = NoHistory(in_dim=4, out_dim=6)
+    B, K = 2, 3
+    has_prev_a = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
+    has_prev_b = torch.zeros(B, K, dtype=torch.bool)
+    feat_a = torch.randn(B, K, 4)
+    feat_b = torch.randn(B, K, 4) * 100
+    out_a = hist(feat_a, has_prev_a)
+    out_b = hist(feat_b, has_prev_b)
+    assert torch.equal(out_a, torch.zeros(B, K, 6))
+    assert torch.equal(out_a, out_b)
+
+
+def test_no_history_uses_base_class_o1_defaults():
+    hist = NoHistory(in_dim=4, out_dim=6)
+    assert hist.init_cache() is None
+    feat = torch.randn(2, 1, 4)
+    has_prev = torch.ones(2, 1, dtype=torch.bool)
+    out, cache = hist.step(feat, has_prev, "unused-cache")
+    assert torch.equal(out, torch.zeros(2, 1, 6))
+    assert cache == "unused-cache"
+
+
+# --- HISTORY_REGISTRY / build_history (gitea #35) ----------------------------
+
+
+def test_history_registry_has_exactly_the_known_histories():
+    assert set(HISTORY_REGISTRY) == {"markov", "attention", "none"}
+
+
+def test_build_history_returns_correct_concrete_type():
+    assert isinstance(build_history("markov", 4, 6), MarkovHistory)
+    assert isinstance(build_history("attention", 4, 8), AttentionHistory)
+    assert isinstance(build_history("none", 4, 6), NoHistory)
+
+
+def test_build_history_unknown_name_raises():
+    with pytest.raises(ValueError):
+        build_history("bogus", 4, 6)
+
+
+def test_build_history_filters_kwargs_by_signature():
+    """Attention-only kwargs (n_heads/n_layers) must be silently dropped when
+    building a MarkovHistory, matching build_router's documented behavior for
+    per-type hyperparameters coexisting in one config."""
+    hist = build_history("markov", 4, 6, n_heads=2, n_layers=1)
+    assert isinstance(hist, MarkovHistory)
+
+
+def test_history_encoder_base_default_init_cache_and_step():
+    """A HistoryEncoder subclass implementing only forward() must still get
+    working O(1) init_cache/step defaults from the base class."""
+
+    class _StubHistory(HistoryEncoder):
+        def forward(self, feat, has_prev):
+            return feat * 2
+
+    hist = _StubHistory()
+    assert hist.init_cache() is None
+    feat = torch.randn(2, 1, 4)
+    has_prev = torch.ones(2, 1, dtype=torch.bool)
+    out, cache = hist.step(feat, has_prev, "unused-cache")
+    assert torch.equal(out, hist.forward(feat, has_prev))
+    assert cache == "unused-cache"
+
+
 # --- Stage2Autoregressive (v0.3.0 step 5) -----------------------------------


@@ -390,10 +557,9 @@ def _build_stage2_ar(
    k_max: int = 5,
    history: str = "markov",
 ) -> Stage2Autoregressive:
-    particle_cfg = {"type": "physical", "emb_dim": emb_dim, "n_layers": 1}
+    particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=emb_dim, n_layers=1)
    if target == "embedding":
-        particle_cfg = dict(particle_cfg)
-        particle_cfg["type"] = "embedding"
+        particle_cfg = gconfig.ConditioningAxisConfig(type="embedding", emb_dim=emb_dim, n_layers=1)
    return Stage2Autoregressive(
        pdg_vocab=5,
        mat_vocab=3,
@@ -405,7 +571,7 @@ def _build_stage2_ar(
        context_dim=8,
        generator=generator,
        k_max=k_max,
-        particle_type_cfg={"target": target, "lambda": 1.0},
+        particle_type_cfg=gconfig.ParticleTypeConfig(target=target),
        history=history,
    )

@@ -423,18 +589,72 @@ def test_stage2_autoregressive_history_invalid_raises():
        _build_stage2_ar("onehot", "wgan", history="bogus")


+def test_stage2_autoregressive_particle_type_n_classes_overrides_conditioning_emb_dim():
+    """gitea #29, Stage2Autoregressive side — see the Stage2OneShot version
+    of this test for the full rationale."""
+    particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=6, n_layers=1)
+    particle_type_cfg = gconfig.ParticleTypeConfig(target="onehot", n_classes=20)
+    model = Stage2Autoregressive(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=particle_cfg,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        cond_out_dim=16,
+        context_dim=8,
+        generator="flow",
+        k_max=5,
+        particle_type_cfg=particle_type_cfg,
+    )
+    assert model.type_dim == 20  # not particle_cfg.emb_dim == 6
+    assert model.type_head is not None
+    assert model.type_head[-1].out_features == 20
+
+
+def test_stage2_autoregressive_n_sec_head_and_type_head_cfg_control_hidden_width_and_depth():
+    """gitea #36, Stage2Autoregressive side — see the Stage2OneShot version
+    of this test for the full rationale."""
+    particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=6, n_layers=1)
+    particle_type_cfg = gconfig.ParticleTypeConfig(target="onehot")
+    model = Stage2Autoregressive(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=particle_cfg,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=40,
+        n_res_blocks=1,
+        cond_out_dim=12,
+        context_dim=8,
+        generator="flow",
+        k_max=5,
+        particle_type_cfg=particle_type_cfg,
+        n_sec_head_cfg={"hidden_ratio": 0.25, "depth": 1},
+        type_head_cfg={"hidden_ratio": 0.75, "depth": 2},
+    )
+    assert model.n_sec_head is not None
+    assert len(model.n_sec_head) == 1
+    assert model.n_sec_head[0].in_features == 12
+    assert model.n_sec_head[0].out_features == 6  # k_max + 1
+
+    assert model.type_head is not None
+    assert len(model.type_head) == 3
+    assert model.type_head[0].out_features == 30  # round(40 * 0.75)
+    assert model.type_head[2].out_features == model.type_dim
+
+
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
@pytest.mark.parametrize("generator", ["wgan", "flow"])
-@pytest.mark.parametrize("history", ["markov", "attention"])
+@pytest.mark.parametrize("history", ["markov", "attention", "none"])
 def test_stage2_autoregressive_forward_shape(target, generator, history):
    B, K, emb_dim = 4, 5, 6
    model = _build_stage2_ar(target, generator, emb_dim=emb_dim, k_max=K, history=history)
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    stage1_out = torch.randn(B, 9)
-    type_dim = stage2_type_dim({"target": target}, emb_dim)
+    type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target=target), emb_dim)
    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
-    token_dim = stage2_trunk_sec_dim({"target": target}, generator, 1, emb_dim)
+    token_dim = stage2_trunk_sec_dim(gconfig.ParticleTypeConfig(target=target), generator, 1, emb_dim)
    if generator == "wgan":
        x_t = torch.randn(B, K, model.noise_dim)
        t = None
@@ -471,7 +691,7 @@ def test_stage2_autoregressive_predict_type_shape():
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    stage1_out = torch.randn(B, 9)
-    type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
+    type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target="onehot"), emb_dim)
    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
    out = model.predict_type(
        cond_cont,
@@ -492,7 +712,7 @@ def test_stage2_autoregressive_predict_type_raises_when_no_type_head(target, gen
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    stage1_out = torch.randn(B, 9)
-    type_dim = stage2_type_dim({"target": target}, emb_dim)
+    type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target=target), emb_dim)
    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
    with pytest.raises(RuntimeError):
        model.predict_type(
@@ -512,7 +732,7 @@ def test_stage2_autoregressive_gradients_flow_wgan_onehot():
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    stage1_out = torch.randn(B, 9)
-    type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
+    type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target="onehot"), emb_dim)
    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
    z = torch.randn(B, K, model.noise_dim)
    gen_out = model(
@@ -537,9 +757,9 @@ def test_stage2_autoregressive_gradients_flow_onehot():
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.zeros(B, 2, dtype=torch.long)
    stage1_out = torch.randn(B, 9)
-    type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
+    type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target="onehot"), emb_dim)
    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
-    token_dim = stage2_trunk_sec_dim({"target": "onehot"}, "flow", 1, emb_dim)
+    token_dim = stage2_trunk_sec_dim(gconfig.ParticleTypeConfig(target="onehot"), "flow", 1, emb_dim)
    x_t = torch.randn(B, K, token_dim)
    t = torch.rand(B, K)
    flow_out = model(
@@ -579,7 +799,7 @@ def test_stage2_autoregressive_history_step_matches_parallel_history_encoder():
    B, K, emb_dim = 3, 6, 6
    model = _build_stage2_ar("physical", "wgan", emb_dim=emb_dim, k_max=K, history="attention")
    model.eval()
-    type_dim = stage2_type_dim({"target": "physical"}, emb_dim)
+    type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target="physical"), emb_dim)
    hist_in_dim = CONT_SLOT_DIM + type_dim
    own_feat = torch.randn(B, K, hist_in_dim)  # token i's own raw feature
    has_prev_full = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
@@ -655,6 +875,81 @@ def test_build_models_share_stages_true_shared_params_are_in_both_stage_paramete
    assert shared_ids <= {id(p) for p in stage2.parameters()}


+def test_condition_encoder_stores_the_exact_particle_and_material_cfg_instances_passed_in():
+    """gitea #38: ConditionEncoder must not round-trip particle_cfg/
+    material_cfg through a dict — the exact ConditioningAxisConfig instance
+    passed in is what `.particle_cfg`/`.material_cfg` hold afterward."""
+    particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
+    material_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
+    enc = ConditionEncoder(pdg_vocab=3, mat_vocab=2, particle_cfg=particle_cfg, material_cfg=material_cfg)
+    assert enc.particle_cfg is particle_cfg
+    assert enc.material_cfg is material_cfg
+
+
+def test_stagemodel_stores_the_exact_particle_type_cfg_instance_passed_in():
+    """gitea #38: a StageModel subclass must not round-trip particle_type_cfg
+    through a dict — the exact ParticleTypeConfig instance passed in is what
+    `.particle_type_cfg` holds afterward."""
+    particle_type_cfg = gconfig.ParticleTypeConfig(target="onehot", n_classes=11)
+    sec_dim = stage2_trunk_sec_dim(particle_type_cfg, "flow", 5, 11)
+    model = Stage2OneShot(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        k_max=5,
+        sec_dim=sec_dim,
+        particle_type_cfg=particle_type_cfg,
+    )
+    assert model.particle_type_cfg is particle_type_cfg
+
+
+def test_build_models_particle_type_cfg_and_conditioning_axes_are_dataclasses_not_dicts():
+    """gitea #38: build_models must pass the parsed ConditioningAxisConfig/
+    ParticleTypeConfig dataclasses themselves down to the model constructors,
+    not re-serialize them to a dict first (the inversion the issue names) —
+    before the fix, .particle_type_cfg was a plain dict (s2_spec.particle_type
+    .to_dict()) and .cond_enc.particle_cfg came from the raw, unparsed
+    conditioning["particle"] dict."""
+    cfg = _minimal_model_config(share_stages=False)
+    cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0, "n_classes": 11}
+    built = build_models(cfg)
+    stage1, stage2 = built["stage1"], built["stage2"]
+    assert stage1 is not None and stage2 is not None
+    assert isinstance(stage2.particle_type_cfg, gconfig.ParticleTypeConfig)
+    assert isinstance(stage1.cond_enc.particle_cfg, gconfig.ConditioningAxisConfig)
+    assert isinstance(stage1.cond_enc.material_cfg, gconfig.ConditioningAxisConfig)
+
+
+def test_build_models_particle_type_n_classes_overrides_conditioning_emb_dim():
+    """gitea #29 end-to-end through build_models: setting
+    stage2_model.particle_type.n_classes independently of
+    conditioning.particle.emb_dim actually resizes the built stage2 model,
+    not just the two lower-level unit tests above."""
+    cfg = _minimal_model_config(share_stages=False)  # conditioning.particle.emb_dim = 4
+    cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0, "n_classes": 11}
+    built = build_models(cfg)
+    assert built["stage2"] is not None
+    assert built["stage2"].type_dim == 11
+
+
+def test_build_critics_particle_type_n_classes_overrides_conditioning_emb_dim():
+    cfg = _minimal_model_config(share_stages=False)  # conditioning.particle.emb_dim = 4
+    cfg["stage2_model"]["generator"] = "wgan"
+    cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0, "n_classes": 4}
+    default_n_classes_critic = build_critics(cfg)["stage2"]
+    assert default_n_classes_critic is not None
+
+    cfg["stage2_model"]["particle_type"]["n_classes"] = 11
+    wider_critic = build_critics(cfg)["stage2"]
+    assert wider_critic is not None
+    # k_max=3 slots, each CONT_SLOT_DIM + n_classes wide under wgan folding —
+    # widening n_classes alone (emb_dim stays 4) must widen the critic input.
+    assert wider_critic.trunk.input_proj.in_features > default_n_classes_critic.trunk.input_proj.in_features
+
+
 # ── build_models/build_critics: DEFAULT_CONFIG fallback drift (issues.md #1) ─


@@ -688,7 +983,40 @@ def _partial_model_config() -> dict:
 def test_build_models_omitted_decoder_and_particle_type_match_default_config():
    built = build_models(_partial_model_config())
    assert isinstance(built["stage2"], Stage2Autoregressive)
-    assert built["stage2"].particle_type_cfg["target"] == "onehot"
+    assert built["stage2"].particle_type_cfg.target == "onehot"
+
+
+def test_build_models_custom_heads_block_controls_head_shapes():
+    """gitea #36: stage{1,2}_model.heads flows all the way from config dict
+    through build_models to the actual constructed head shapes."""
+    cfg = _partial_model_config()
+    cfg["stage1_model"] = {
+        "active": True,
+        "hidden_dim": 40,
+        "n_res_blocks": 1,
+        "heads": {"n_sec": {"hidden_ratio": 0.25, "depth": 1}},
+    }
+    cfg["stage2_model"]["decoder"] = "one_shot"
+    cfg["stage2_model"]["generator"] = "flow"  # wgan folds the type slice; no separate type_head
+    cfg["stage2_model"]["n_sec"] = {"owner": "stage1"}
+    cfg["stage2_model"]["heads"] = {
+        "n_sec": {"hidden_ratio": 0.25, "depth": 1},
+        "type": {"hidden_ratio": 0.75, "depth": 2},
+    }
+
+    built = build_models(cfg)
+    stage1, stage2 = built["stage1"], built["stage2"]
+    assert stage1 is not None
+    assert stage2 is not None
+
+    assert stage1.n_sec_head is not None
+    assert len(stage1.n_sec_head) == 1  # owner=stage1, so stage1 builds it
+    assert stage2.n_sec_head is None  # owner=stage1, so stage2 doesn't
+
+    assert isinstance(stage2, Stage2OneShot)
+    assert stage2.type_head is not None
+    assert len(stage2.type_head) == 3
+    assert stage2.type_head[0].out_features == 6  # round(8 * 0.75)


 def test_build_critics_omitted_particle_type_matches_default_config():
@@ -696,15 +1024,247 @@ def test_build_critics_omitted_particle_type_matches_default_config():
    cfg["stage2_model"]["generator"] = "wgan"
    onehot_critic = build_critics(cfg)["stage2"]
    assert onehot_critic is not None
-    onehot_in_dim = onehot_critic.input_proj.in_features
+    onehot_in_dim = onehot_critic.trunk.input_proj.in_features

    cfg["stage2_model"]["particle_type"] = {"target": "physical"}
    physical_critic = build_critics(cfg)["stage2"]
    assert physical_critic is not None
-    physical_in_dim = physical_critic.input_proj.in_features
+    physical_in_dim = physical_critic.trunk.input_proj.in_features

    # onehot's per-slot type width is emb_dim classes vs. physical's fixed
    # (log-mass, charge) pair — different unless emb_dim happens to be 2, so
    # this also confirms the critic was actually built in onehot mode by
    # default, not silently falling back to physical.
    assert onehot_in_dim != physical_in_dim
+
+
+# ── build_critics: critic_hidden_dim/critic_n_res_blocks honoured (gitea #28) ─
+
+
+def test_build_critics_stage1_critic_hidden_dim_and_n_res_blocks_override_generator_size():
+    cfg = _minimal_model_config(share_stages=False)
+    cfg["stage1_model"]["generator"] = "wgan"
+    cfg["stage1_model"]["hidden_dim"] = 8
+    cfg["stage1_model"]["n_res_blocks"] = 1
+
+    inherited = build_critics(cfg)["stage1"]
+    assert inherited is not None
+    assert inherited.trunk.input_proj.out_features == 8
+    assert len(inherited.trunk.blocks) == 1
+
+    cfg["stage1_model"]["wgan"]["critic_hidden_dim"] = 16
+    cfg["stage1_model"]["wgan"]["critic_n_res_blocks"] = 3
+    overridden = build_critics(cfg)["stage1"]
+    assert overridden is not None
+    assert overridden.trunk.input_proj.out_features == 16
+    assert len(overridden.trunk.blocks) == 3
+
+
+def test_build_critics_stage2_critic_hidden_dim_and_n_res_blocks_override_generator_size():
+    cfg = _minimal_model_config(share_stages=False)
+    cfg["stage2_model"]["generator"] = "wgan"
+    cfg["stage2_model"]["hidden_dim"] = 8
+    cfg["stage2_model"]["n_res_blocks"] = 1
+
+    inherited = build_critics(cfg)["stage2"]
+    assert inherited is not None
+    assert inherited.trunk.input_proj.out_features == 8
+    assert len(inherited.trunk.blocks) == 1
+
+    cfg["stage2_model"]["wgan"]["critic_hidden_dim"] = 16
+    cfg["stage2_model"]["wgan"]["critic_n_res_blocks"] = 3
+    overridden = build_critics(cfg)["stage2"]
+    assert overridden is not None
+    assert overridden.trunk.input_proj.out_features == 16
+    assert len(overridden.trunk.blocks) == 3
+
+
+# ── StageModel base (gitea #39): Stage1Model/Stage2OneShot/Stage2Autoregressive
+# scaffolding — construction order, and therefore fresh-init RNG draw order and
+# state_dict key set, must stay byte-for-byte what it was before the base class
+# existed. ------------------------------------------------------------------
+
+_STAGE_HIDDEN_DIM = 32
+_STAGE_N_BLOCKS = 2
+_STAGE_COND_OUT_DIM = 16
+
+
+def _resblock_keys(prefix: str) -> set[str]:
+    return {
+        f"{prefix}.norm.weight",
+        f"{prefix}.norm.bias",
+        f"{prefix}.linear1.weight",
+        f"{prefix}.linear1.bias",
+        f"{prefix}.cond_proj.weight",
+        f"{prefix}.linear2.weight",
+        f"{prefix}.linear2.bias",
+    }
+
+
+def _trunk_keys(prefix: str = "trunk") -> set[str]:
+    keys = {
+        f"{prefix}.input_proj.weight",
+        f"{prefix}.input_proj.bias",
+        f"{prefix}.out_proj.weight",
+        f"{prefix}.out_proj.bias",
+    }
+    for i in range(_STAGE_N_BLOCKS):
+        keys |= _resblock_keys(f"{prefix}.blocks.{i}")
+    return keys
+
+
+def _cond_enc_keys() -> set[str]:
+    return {
+        "cond_enc.mlp.0.weight",
+        "cond_enc.mlp.0.bias",
+        "cond_enc.mlp.2.weight",
+        "cond_enc.mlp.2.bias",
+        "cond_enc.particle_mlp.0.weight",
+        "cond_enc.particle_mlp.0.bias",
+        "cond_enc.material_mlp.0.weight",
+        "cond_enc.material_mlp.0.bias",
+    }
+
+
+def _fuse_keys(name: str) -> set[str]:
+    return {f"{name}.0.weight", f"{name}.0.bias"}
+
+
+def _head_keys(name: str) -> set[str]:
+    return {f"{name}.0.weight", f"{name}.0.bias", f"{name}.2.weight", f"{name}.2.bias"}
+
+
+def _expected_stage_keys(*, has_time: bool, extra: set[str]) -> set[str]:
+    keys = _cond_enc_keys() | _trunk_keys() | extra
+    if has_time:
+        keys.add("time_emb.freqs")
+    return keys
+
+
+@pytest.mark.parametrize("generator", ["flow", "wgan"])
+def test_stage1_model_state_dict_keys_unchanged_by_stagemodel_refactor(generator):
+    model = Stage1Model(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=_STAGE_HIDDEN_DIM,
+        n_res_blocks=_STAGE_N_BLOCKS,
+        cond_out_dim=_STAGE_COND_OUT_DIM,
+        generator=generator,
+        time_dim=8,
+        noise_dim=8,
+        n_sec_head_k_max=15,
+    )
+    expected = _expected_stage_keys(
+        has_time=build_objective(generator).needs_time,
+        extra=_head_keys("n_sec_head"),
+    )
+    assert set(model.state_dict().keys()) == expected
+
+
+@pytest.mark.parametrize("generator", ["flow", "wgan"])
+def test_stage2_oneshot_state_dict_keys_unchanged_by_stagemodel_refactor(generator):
+    model = Stage2OneShot(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=_STAGE_HIDDEN_DIM,
+        n_res_blocks=_STAGE_N_BLOCKS,
+        cond_out_dim=_STAGE_COND_OUT_DIM,
+        generator=generator,
+        time_dim=8,
+        noise_dim=8,
+        k_max=15,
+    )
+    extra = _head_keys("n_sec_head") | {"context_adapter.proj.weight", "context_adapter.proj.bias"} | _fuse_keys("fuse")
+    expected = _expected_stage_keys(has_time=build_objective(generator).needs_time, extra=extra)
+    assert set(model.state_dict().keys()) == expected
+
+
+@pytest.mark.parametrize("generator", ["flow", "wgan"])
+def test_stage2_autoregressive_state_dict_keys_unchanged_by_stagemodel_refactor(generator):
+    model = Stage2Autoregressive(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=_STAGE_HIDDEN_DIM,
+        n_res_blocks=_STAGE_N_BLOCKS,
+        cond_out_dim=_STAGE_COND_OUT_DIM,
+        generator=generator,
+        time_dim=8,
+        noise_dim=8,
+        k_max=15,
+    )
+    extra = (
+        _head_keys("n_sec_head")
+        | {"context_adapter.proj.weight", "context_adapter.proj.bias"}
+        | _fuse_keys("base_fuse")
+        | _fuse_keys("token_fuse")
+        | {"history_encoder.start", "history_encoder.mlp.0.weight", "history_encoder.mlp.0.bias"}
+    )
+    expected = _expected_stage_keys(has_time=build_objective(generator).needs_time, extra=extra)
+    assert set(model.state_dict().keys()) == expected
+
+
+@pytest.mark.parametrize("cls", [Stage1Model, Stage2OneShot, Stage2Autoregressive])
+def test_stage_classes_are_stagemodel_subclasses(cls):
+    assert issubclass(cls, StageModel)
+
+
+@pytest.mark.parametrize("cls", [Stage1Model, Stage2OneShot, Stage2Autoregressive])
+@pytest.mark.parametrize("generator", ["flow", "ddpm", "wgan"])
+def test_stagemodel_time_emb_matches_objective_needs_time(cls, generator):
+    kwargs = dict(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=_STAGE_HIDDEN_DIM,
+        n_res_blocks=_STAGE_N_BLOCKS,
+        cond_out_dim=_STAGE_COND_OUT_DIM,
+        generator=generator,
+        time_dim=8,
+        noise_dim=8,
+    )
+    if cls is Stage1Model:
+        kwargs["n_sec_head_k_max"] = 15
+    else:
+        kwargs["k_max"] = 15
+    model = cls(**kwargs)
+    assert model.generator_kind == generator
+    assert model.noise_dim == 8
+    assert (model.time_emb is not None) == build_objective(generator).needs_time
+
+
+# ── CriticModel uses the trunk/block registries + StageModel base (gitea #57) ─
+
+
+def test_critic_model_is_stagemodel_subclass():
+    assert issubclass(CriticModel, StageModel)
+
+
+@pytest.mark.parametrize("stage", ["stage1", "stage2"])
+def test_build_critics_threads_trunk_type_from_generator_config(stage):
+    cfg = _minimal_model_config(share_stages=False)
+    cfg["stage1_model"]["generator"] = "wgan"
+    cfg["stage2_model"]["generator"] = "wgan"
+    cfg[f"{stage}_model"]["trunk"] = {"type": "linear"}
+
+    critic = build_critics(cfg)[stage]
+    assert critic is not None
+    assert isinstance(critic.trunk, LinearTrunk)
+
+
+@pytest.mark.parametrize("stage", ["stage1", "stage2"])
+def test_build_critics_threads_block_conditioning_from_generator_config(stage):
+    cfg = _minimal_model_config(share_stages=False)
+    cfg["stage1_model"]["generator"] = "wgan"
+    cfg["stage2_model"]["generator"] = "wgan"
+    cfg[f"{stage}_model"]["trunk"] = {"block_conditioning": "film"}
+
+    critic = build_critics(cfg)[stage]
+    assert critic is not None
+    assert all(isinstance(block, FilmResBlock) for block in critic.trunk.blocks)
@@ -0,0 +1,251 @@
+"""Tests for `giant/model/objectives.py` — the generator/objective registry
+(gitea #32) that replaced bare `generator in ("flow", "ddpm", "wgan")`
+string checks scattered across models.py/sample.py/builders.py/
+stage2_inputs.py/trainers.py."""
+
+import pytest
+import torch
+
+from giant.config import ConditioningAxisConfig
+from giant.constants import COND_DIM, CONT_SLOT_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM
+from giant.model.network import (
+    OBJECTIVE_REGISTRY,
+    DdpmObjective,
+    FlowObjective,
+    Stage1Model,
+    Stage2Autoregressive,
+    Stage2OneShot,
+    WganObjective,
+    build_objective,
+)
+from giant.model.schedule import CosineSchedule, flow_matching_loss, flow_matching_loss_secondary
+
+_PHYS_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
+
+
+def _cond(B: int, pdg: int = 3, mat: int = 2) -> tuple[torch.Tensor, torch.Tensor]:
+    cond_cont = torch.randn(B, COND_DIM)
+    cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
+    return cond_cont, cond_cat
+
+
+# ── registry ─────────────────────────────────────────────────────────────
+
+
+def test_registry_has_exactly_the_three_known_objectives():
+    assert set(OBJECTIVE_REGISTRY) == {"flow", "ddpm", "wgan"}
+
+
+def test_build_objective_returns_correct_concrete_type():
+    assert isinstance(build_objective("flow"), FlowObjective)
+    assert isinstance(build_objective("ddpm"), DdpmObjective)
+    assert isinstance(build_objective("wgan"), WganObjective)
+
+
+def test_build_objective_unknown_name_raises():
+    with pytest.raises(ValueError, match="unknown generator/objective"):
+        build_objective("bogus")
+
+
+def test_build_objective_filters_kwargs_by_signature():
+    # FlowObjective takes no constructor args — n_steps (a DdpmObjective-only
+    # kwarg) must be silently dropped, not raise a TypeError.
+    build_objective("flow", n_steps=500)
+    ddpm = build_objective("ddpm", n_steps=250)
+    assert isinstance(ddpm, DdpmObjective)
+    assert ddpm.n_steps == 250
+
+
+# ── flags ────────────────────────────────────────────────────────────────
+
+
+def test_flow_objective_flags():
+    obj = build_objective("flow")
+    assert obj.needs_time is True
+    assert obj.is_adversarial is False
+    assert obj.folds_type_slice is False
+    assert obj.supports_stage2_decoder is True
+
+
+def test_ddpm_objective_flags():
+    obj = build_objective("ddpm")
+    assert obj.needs_time is True
+    assert obj.is_adversarial is False
+    assert obj.folds_type_slice is False
+    assert obj.supports_stage2_decoder is False
+
+
+def test_wgan_objective_flags():
+    obj = build_objective("wgan")
+    assert obj.needs_time is False
+    assert obj.is_adversarial is True
+    assert obj.folds_type_slice is True
+    assert obj.supports_stage2_decoder is True
+
+
+# ── trunk_in_dim ─────────────────────────────────────────────────────────
+
+
+def test_trunk_in_dim_flow_and_ddpm_pass_through_out_dim():
+    assert build_objective("flow").trunk_in_dim(out_dim=9, noise_dim=8) == 9
+    assert build_objective("ddpm").trunk_in_dim(out_dim=9, noise_dim=8) == 9
+
+
+def test_trunk_in_dim_wgan_uses_noise_dim():
+    assert build_objective("wgan").trunk_in_dim(out_dim=9, noise_dim=8) == 8
+
+
+# ── ddpm schedule ────────────────────────────────────────────────────────
+
+
+def test_ddpm_build_schedule_has_requested_length():
+    schedule = build_objective("ddpm").build_schedule(n_steps=17, device=torch.device("cpu"))
+    assert isinstance(schedule, CosineSchedule)
+    assert schedule.T == 17
+
+
+def test_flow_and_wgan_build_schedule_is_none():
+    assert build_objective("flow").build_schedule(100, torch.device("cpu")) is None
+    assert build_objective("wgan").build_schedule(100, torch.device("cpu")) is None
+
+
+# ── stage1_loss parity ──────────────────────────────────────────────────
+
+
+def test_flow_objective_stage1_loss_matches_direct_call():
+    torch.manual_seed(0)
+    model = Stage1Model(
+        pdg_vocab=3, mat_vocab=2, particle_cfg=_PHYS_CFG, material_cfg=_PHYS_CFG, hidden_dim=16, n_res_blocks=1
+    )
+    cond_cont, cond_cat = _cond(4)
+    x1 = torch.randn(4, X_DIM)
+
+    torch.manual_seed(1)
+    expected = flow_matching_loss(model, x1, cond_cont, cond_cat)
+    torch.manual_seed(1)
+    actual = build_objective("flow").stage1_loss(model, x1, cond_cont, cond_cat)
+    assert torch.allclose(actual, expected)
+
+
+def test_ddpm_objective_stage1_loss_matches_direct_call():
+    torch.manual_seed(0)
+    model = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=_PHYS_CFG,
+        material_cfg=_PHYS_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        generator="ddpm",
+    )
+    cond_cont, cond_cat = _cond(4)
+    x1 = torch.randn(4, X_DIM)
+    objective = build_objective("ddpm", n_steps=50)
+    schedule = objective.build_schedule(50, torch.device("cpu"))
+    assert isinstance(schedule, CosineSchedule)
+
+    torch.manual_seed(1)
+    expected = schedule.loss(model, x1, cond_cont, cond_cat)
+    torch.manual_seed(1)
+    actual = objective.stage1_loss(model, x1, cond_cont, cond_cat, schedule=schedule)
+    assert torch.allclose(actual, expected)
+
+
+def test_ddpm_objective_stage1_loss_requires_a_schedule():
+    model = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=_PHYS_CFG,
+        material_cfg=_PHYS_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        generator="ddpm",
+    )
+    cond_cont, cond_cat = _cond(4)
+    with pytest.raises(AssertionError):
+        build_objective("ddpm").stage1_loss(model, torch.randn(4, X_DIM), cond_cont, cond_cat, schedule=None)
+
+
+# ── stage2_loss dispatch ─────────────────────────────────────────────────
+
+
+def test_flow_objective_stage2_loss_one_shot_matches_direct_call():
+    torch.manual_seed(0)
+    B, k_max = 4, 5
+    sec_dim = k_max * SEC_SLOT_DIM
+    model = Stage2OneShot(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=_PHYS_CFG,
+        material_cfg=_PHYS_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        generator="flow",
+        sec_dim=sec_dim,
+        k_max=k_max,
+    )
+    cond_cont, cond_cat = _cond(B)
+    stage1_ctx = torch.randn(B, X_DIM)
+    x1_s2 = torch.randn(B, sec_dim)
+    sec_mask = torch.ones(B, k_max, dtype=torch.bool)
+
+    torch.manual_seed(1)
+    expected = flow_matching_loss_secondary(model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=None)
+    torch.manual_seed(1)
+    actual = build_objective("flow").stage2_loss(
+        model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=None, ar_inputs=None
+    )
+    assert torch.allclose(actual, expected)
+
+
+def test_flow_objective_stage2_loss_dispatches_to_ar_when_ar_inputs_given():
+    torch.manual_seed(0)
+    B, k_max = 4, 5
+    model = Stage2Autoregressive(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=_PHYS_CFG,
+        material_cfg=_PHYS_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        generator="flow",
+        k_max=k_max,
+    )
+    cond_cont, cond_cat = _cond(B)
+    stage1_ctx = torch.randn(B, X_DIM)
+    token_dim = CONT_SLOT_DIM + PARTICLE_PHYS_DIM
+    x1_s2 = torch.randn(B, k_max, token_dim)
+    sec_mask = torch.ones(B, k_max, dtype=torch.bool)
+    ar_inputs = {
+        "history_feat": torch.randn(B, k_max, token_dim),
+        "has_prev": torch.ones(B, k_max, dtype=torch.bool),
+        "remaining_frac": torch.rand(B, k_max),
+        "slot_idx": torch.linspace(0, 1, k_max).unsqueeze(0).expand(B, -1),
+    }
+
+    loss = build_objective("flow").stage2_loss(
+        model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=None, ar_inputs=ar_inputs
+    )
+    assert loss.dim() == 0
+    assert torch.isfinite(loss)
+
+
+def test_ddpm_objective_stage2_loss_not_implemented():
+    dummy_model = torch.nn.Module()
+    dummy = torch.zeros(1)
+    with pytest.raises(NotImplementedError):
+        build_objective("ddpm").stage2_loss(
+            dummy_model, dummy, dummy, dummy, dummy, torch.ones(1, 1, dtype=torch.bool), type_dim=None
+        )
+
+
+def test_wgan_objective_has_no_loss_methods():
+    dummy_model = torch.nn.Module()
+    dummy = torch.zeros(1)
+    objective = build_objective("wgan")
+    with pytest.raises(NotImplementedError):
+        objective.stage1_loss(dummy_model, dummy, dummy, dummy)
+    with pytest.raises(NotImplementedError):
+        objective.stage2_loss(
+            dummy_model, dummy, dummy, dummy, dummy, torch.ones(1, 1, dtype=torch.bool), type_dim=None
+        )
@@ -4,6 +4,7 @@ import numpy as np
 import pytest
 import torch

+from giant.config import ConditioningAxisConfig
 from giant.constants import (
    COND_DIM,
    CONT_SLOT_DIM,
@@ -23,9 +24,9 @@ from giant.sample import sample_secondaries
 # ── helpers ──────────────────────────────────────────────────────────────────


-def _particle_material_cfg(conditioning: str) -> tuple[dict, dict]:
-    cfg = {"type": conditioning, "emb_dim": 16, "n_layers": 1}
-    return dict(cfg), dict(cfg)
+def _particle_material_cfg(conditioning: str) -> tuple[ConditioningAxisConfig, ConditioningAxisConfig]:
+    cfg = ConditioningAxisConfig(type=conditioning, emb_dim=16, n_layers=1)
+    return cfg, cfg


 def _stage1(pdg=3, mat=2, conditioning="embedding"):
@@ -626,3 +627,53 @@ def test_decode_secondaries_mass_charge_round_trip_with_normalizer():
    _, _, sec_mass, sec_charge, _ = decode_secondaries(sec_cont_normed, n_sec, e_sec, pre_dir, sec_phys_normalizer=norm)
    assert sec_mass[0, 0] == pytest.approx(938.27208943, abs=1e-2)
    assert sec_charge[0, 0] == pytest.approx(1.0, abs=1e-4)
+
+
+def test_decode_secondaries_extreme_negative_log_mass_stays_nonnegative():
+    from giant.data.transforms import decode_secondaries, log_transform
+
+    N = 1
+    e_sec = np.array([5.0], dtype=np.float32)
+    n_sec = np.array([1])
+    pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
+
+    sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32)
+    sec_cont[0, 0, 0] = 10.0  # stick logit -> ~all of e_sec
+    sec_cont[0, 0, 1:4] = [0, 0, 1]
+    sec_cont[0, 0, 4] = -50.0  # raw model prediction: extremely negative log_mass
+    sec_cont[0, 0, 5] = 1.0
+
+    _, _, sec_mass, _, _ = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
+
+    # A raw model prediction isn't itself the output of log_transform, so
+    # naively applying inv_log_transform can undershoot zero (see
+    # decode_secondaries) — which then crashes the next log_transform call
+    # once this mass is fed back in as conditioning during rollout. The
+    # float32 residual from clipping can land a hair below zero, but must
+    # stay well above -eps so log_transform(mass) stays finite.
+    assert sec_mass[0, 0] > -1e-8
+    log_transform(sec_mass[0, 0])
+
+
+def test_decode_secondaries_extreme_positive_log_mass_stays_finite():
+    from giant.data.transforms import decode_secondaries, log_transform
+
+    N = 1
+    e_sec = np.array([5.0], dtype=np.float32)
+    n_sec = np.array([1])
+    pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
+
+    sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32)
+    sec_cont[0, 0, 0] = 10.0  # stick logit -> ~all of e_sec
+    sec_cont[0, 0, 1:4] = [0, 0, 1]
+    sec_cont[0, 0, 4] = 200.0  # raw model prediction: extremely positive log_mass
+    sec_cont[0, 0, 5] = 1.0
+
+    _, _, sec_mass, _, _ = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
+
+    # Mirror image of the extreme-negative case above: exp(log_mass)
+    # overflows float32 to inf for an unclipped raw prediction this large,
+    # which then crashes the next log_transform call the same way a
+    # negative mass would.
+    assert np.isfinite(sec_mass[0, 0])
+    log_transform(sec_mass[0, 0])
@@ -152,28 +152,60 @@ def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch):
    assert "normalizer: cache hit" in joined


-def test_run_train_job_builds_caches_and_persists_pdg_topn_map(tmp_path, data):
+def test_run_train_job_builds_caches_and_persists_sec_type_topn_map(tmp_path, data):
    """DEFAULT_CONFIG's stage2_model.particle_type.target defaults to
-    "onehot" — a plain _tiny_cfg() run must
-    build the shared pdg top-N map, cache it in the setup-cache sidecar, and
-    persist it into the checkpoint, with no extra config needed."""
+    "onehot" while conditioning.particle.type stays "physical" — a plain
+    _tiny_cfg() run must build the secondary-type-only pdg top-N map (gitea
+    #29: no longer shared with any conditioning-side onehot map), cache it in
+    the setup-cache sidecar, and persist it into the checkpoint's
+    sec_type_topn_map key, with no extra config needed. pdg_topn_map
+    (conditioning-only) stays unbuilt since conditioning.particle.type is
+    "physical" here."""
    echo1 = _run(data, tmp_path / "out1")
    assert any("building pdg top-N map" in m for m in echo1)

    loaded = setup_cache.load(data, [data])
    assert loaded is not None
-    key = setup_cache.topn_key("pdg", 4)  # conditioning.particle.emb_dim = 4
+    # stage2_model.particle_type.n_classes = 0 -> conditioning.particle.emb_dim = 4
+    key = setup_cache.topn_key("pdg", 4)
    assert key in loaded.topn_maps
    assert set(loaded.topn_maps[key].class_map.keys()) >= {11, 22}

    ckpt = torch.load(tmp_path / "out1" / "last.pt", weights_only=False)
-    assert "pdg_topn_map" in ckpt
-    assert set(ckpt["pdg_topn_map"]["class_map"].keys()) >= {"11", "22"}
+    assert ckpt.get("pdg_topn_map") is None
+    assert "sec_type_topn_map" in ckpt
+    assert set(ckpt["sec_type_topn_map"]["class_map"].keys()) >= {"11", "22"}

    echo2 = _run(data, tmp_path / "out2")
    assert any("pdg top-N map: cache hit" in m for m in echo2)


+def test_run_train_job_independent_cond_and_sec_type_topn_maps(tmp_path, data):
+    """conditioning.particle.type="onehot" and
+    stage2_model.particle_type.target="onehot" with different class counts
+    (gitea #29's fix: stage2_model.particle_type.n_classes decouples the two)
+    build two distinct top-N maps, cached under their own (axis, n_classes)
+    key and persisted under two distinct checkpoint keys — no longer forced
+    to share conditioning.particle.emb_dim."""
+    cfg = _tiny_cfg()
+    cfg["conditioning"]["particle"]["type"] = "onehot"  # emb_dim = 4, from _tiny_cfg
+    cfg["stage2_model"]["particle_type"]["n_classes"] = 3
+    echo = _run(data, tmp_path / "out", cfg=cfg)
+    assert any("mapped to 4 classes" in m for m in echo)
+    assert any("mapped to 3 classes" in m for m in echo)
+
+    loaded = setup_cache.load(data, [data])
+    assert loaded is not None
+    cond_key = setup_cache.topn_key("pdg", 4)
+    type_key = setup_cache.topn_key("pdg", 3)
+    assert cond_key in loaded.topn_maps
+    assert type_key in loaded.topn_maps
+
+    ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
+    assert ckpt.get("pdg_topn_map") is not None
+    assert ckpt.get("sec_type_topn_map") is not None
+
+
 def test_run_train_job_builds_caches_and_persists_material_topn_map(tmp_path, data):
    """conditioning.material.type="onehot" is an independent axis from the
    pdg one above, with its own build/cache-hit branch in run_setup_stage —
@@ -199,7 +231,7 @@ def test_run_train_job_builds_caches_and_persists_material_topn_map(tmp_path, da

 def test_run_train_job_no_topn_map_for_physical_target(tmp_path, data):
    cfg = _tiny_cfg()
-    cfg["stage2_model"]["particle_type"] = {"target": "physical", "lambda": 1.0}
+    cfg["stage2_model"]["particle_type"].update({"target": "physical", "lambda": 1.0})
    echo = _run(data, tmp_path / "out", cfg=cfg)
    assert not any("top-N map" in m for m in echo)

@@ -0,0 +1,77 @@
+"""Config-correctness tests for the CI version-bump/tag/changelog automation
+(gitea #50). The workflow YAML itself can only be exercised by a real push to
+master, so these check the two config files it drives (.bumpversion.toml,
+cliff.toml) against real repo content instead.
+"""
+
+import re
+import shutil
+import subprocess
+import tomllib
+from pathlib import Path
+
+import pytest
+
+_ROOT = Path(__file__).resolve().parents[1]
+
+
+def test_bumpversion_search_pattern_matches_pyproject():
+    bump_config = tomllib.loads((_ROOT / ".bumpversion.toml").read_text())["tool"]["bumpversion"]
+    current_version = bump_config["current_version"]
+    search = bump_config["files"][0]["search"].format(current_version=current_version)
+
+    pyproject = (_ROOT / "pyproject.toml").read_text()
+    assert search in pyproject, (
+        f"bumpversion search pattern {search!r} (rendered from .bumpversion.toml's "
+        f"current_version={current_version!r}) not found in pyproject.toml — "
+        "the bump would silently edit nothing"
+    )
+
+
+@pytest.mark.skipif(shutil.which("git-cliff") is None, reason="git-cliff binary not on PATH")
+def test_cliff_config_groups_and_links_commits(tmp_path):
+    repo = tmp_path / "repo"
+    repo.mkdir()
+    subprocess.run(["git", "init", "-q"], cwd=repo, check=True)
+    subprocess.run(["git", "config", "user.name", "test"], cwd=repo, check=True)
+    subprocess.run(["git", "config", "user.email", "test@example.com"], cwd=repo, check=True)
+
+    _commit(repo, "Add class-balanced secondary particle-type loss (gitea #44)")
+    _commit(repo, "Fix leaking secondary energy budget")
+    _commit(repo, "Merge pull request 'Add X' (#1) from fix/issue-1 into master")
+    _commit(repo, "chore: bump version 0.3.3 -> 0.3.4 [skip ci]")
+
+    result = subprocess.run(
+        [
+            "git-cliff",
+            "--config",
+            str(_ROOT / "cliff.toml"),
+            "--repository",
+            str(repo),
+            "--tag",
+            "v0.3.4",
+            "--unreleased",
+        ],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    changelog = result.stdout
+
+    assert "## [0.3.4]" in changelog
+    assert "### Added" in changelog
+    assert "### Fixed" in changelog
+    assert re.search(
+        r"\[gitea #44\]\(https://git\.larsbogner\.de/lars/giant/issues/44\)",
+        changelog,
+    )
+    assert "Add class-balanced secondary particle-type loss" in changelog
+    assert "Fix leaking secondary energy budget" in changelog
+    assert "Merge pull request" not in changelog
+    assert "skip ci" not in changelog
+
+
+def _commit(repo: Path, message: str) -> None:
+    (repo / "f.txt").write_text(message)
+    subprocess.run(["git", "add", "f.txt"], cwd=repo, check=True)
+    subprocess.run(["git", "commit", "-q", "-m", message], cwd=repo, check=True)
--- a/Show More
+++ b/Show More