Merge pull request 'V0.3.0 stage2 autoregressive' (#27) from v0.3.0-stage2-autoregressive into master

Reviewed-on: #27

CLAUDE.md

@@ -22,7 +22,7 @@ giant analyze render <run_dir> --gallery                   # render PDFs + HTML
 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.
@@ -91,6 +91,6 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep
 
 A sampling-calorimeter (multi-material) dataset is still a planned future direction, not yet built. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`).
 
-**v0.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.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. This config break is why v0.2-shaped configs/checkpoints need migrating at all (`config.migrate_config`, `model.network._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.

README.md

@@ -93,7 +93,7 @@ giant/
 │   │   ├── condor.py       #   prep / compute-one / submit-description 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`)
+├── 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

giant/_migration.py

@@ -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."
+    )

giant/analysis/router_gating.py

@@ -66,27 +66,11 @@ class _RouterHandle:
     router_type: str
 
 
-def _conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]:
-    """(particle_conditioning, material_conditioning) for
-    `giant.data.transforms.build_cond_features` — from either a v0.2
-    checkpoint's flat `model_config["conditioning"]` (one shared string, same
-    for both axes) or a new-format one (independent
-    `model_config["conditioning"]["particle"/"material"]["type"]` — the two
-    axes may differ). Mirrors
-    `giant.cli._conditioning_axes`."""
-    raw = model_cfg.get("conditioning", default)
-    if isinstance(raw, dict):
-        return (
-            raw.get("particle", {}).get("type", default),
-            raw.get("material", {}).get("type", default),
-        )
-    return raw, raw
-
-
 def load_router(checkpoint: str | Path) -> _RouterHandle | None:
     """Load a checkpoint's Stage-1 router, or None if it isn't a MoE checkpoint."""
     import torch
 
+    from giant.checkpoint_io import conditioning_axes
     from giant.data.transforms import Normalizer
     from giant.model.network import build_models
 
@@ -110,7 +94,7 @@ def load_router(checkpoint: str | Path) -> _RouterHandle | None:
     if router is None:
         return None
 
-    particle_conditioning, material_conditioning = _conditioning_axes(model_cfg)
+    particle_conditioning, material_conditioning = conditioning_axes(model_cfg)
     return _RouterHandle(
         router=router,
         pdg_map={int(k): v for k, v in ckpt["pdg_map"].items()},

giant/analysis/runtime_estimate.py

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

giant/checkpoint_io.py

@@ -0,0 +1,234 @@
+"""Load a trained checkpoint into ready-to-run models (giant.cli's `predict`/`rollout`).
+
+Both commands need the same ~15 steps to go from a checkpoint path to two
+`eval()`-mode models plus their normalizers/vocab maps: load the pickle,
+validate it carries what current code expects, resolve which conditioning
+mode each axis was trained with, restore the top-N vocab maps (if the
+checkpoint used one-hot conditioning), rebuild the normalizers, construct the
+model from `model_config`, and load the requested (raw or EMA) weights. This
+used to be duplicated near-verbatim in both commands (issues.md Issue 5) —
+`load_for_inference` is the single implementation.
+
+This module intentionally has no Typer dependency, so it can be unit-tested
+directly and imported from non-CLI code (`giant.analysis.router_gating`,
+lazily — see that module's docstring for why). Failures raise
+`CheckpointCompatibilityError` with the same wording the CLI has always
+shown; the CLI layer catches it and does the `typer.echo`/`Exit(1)`.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+from torch import nn
+
+from giant import config as gconfig
+from giant.constants import K_MAX
+from giant.data.loader import TopNMap
+from giant.data.setup_cache import topnmap_from_json
+from giant.data.transforms import Normalizer
+from giant.model.network import build_models
+
+
+class CheckpointCompatibilityError(Exception):
+    """Checkpoint is missing something `load_for_inference` needs."""
+
+
+def conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]:
+    """(particle_conditioning, material_conditioning) for
+    `giant.data.transforms.build_cond_features`/`build_features` — from
+    either a v0.2 checkpoint's flat `model_config["conditioning"]` (one
+    shared string, same for both axes) or a new-format one (independent
+    `model_config["conditioning"]["particle"/"material"]["type"]` — the two
+    axes are configured independently and may differ)."""
+    raw = model_cfg.get("conditioning", default)
+    if isinstance(raw, dict):
+        return (
+            raw.get("particle", {}).get("type", default),
+            raw.get("material", {}).get("type", default),
+        )
+    return raw, raw
+
+
+def stage_cfg(model_cfg: dict, stage: str) -> dict:
+    """`model_cfg[f"{stage}_model"]` for a new-format model_config, `{}` for
+    a v0.2 flat one (whose ddpm schedule always used `CosineSchedule`'s own
+    default `T=1000` — never a config key — and which never had
+    `particle_type` at all, so `{}` is the correct fallback for both
+    `ddpm_steps`/`particle_type_other_policy` below)."""
+    val = model_cfg.get(f"{stage}_model")
+    return val if isinstance(val, dict) else {}
+
+
+def ddpm_steps(model_cfg: dict, stage: str) -> int:
+    return stage_cfg(model_cfg, stage).get("ddpm", {}).get("n_steps", 1000)
+
+
+def particle_type_other_policy(model_cfg: dict) -> str:
+    return stage_cfg(model_cfg, "stage2").get("particle_type", {}).get("other_policy", "sample")
+
+
+def load_pdg_topn_map(ckpt: dict) -> TopNMap | None:
+    """`ckpt["pdg_topn_map"]` as a `giant.data.loader.TopNMap`, or `None` if
+    this checkpoint's conditioning/particle_type never needed one (see
+    `giant.pipeline.run_setup_stage`, which only populates it when
+    `conditioning.particle.type` or `stage2_model.particle_type.target` is
+    `"onehot"`)."""
+    raw = ckpt.get("pdg_topn_map")
+    return topnmap_from_json(raw, axis="pdg") if raw is not None else None
+
+
+def load_mat_topn_map(ckpt: dict) -> TopNMap | None:
+    """`ckpt["mat_topn_map"]` as a `giant.data.loader.TopNMap`, or `None` if
+    this checkpoint's `conditioning.material.type` was never `"onehot"` (see
+    `giant.pipeline.run_setup_stage`)."""
+    raw = ckpt.get("mat_topn_map")
+    return topnmap_from_json(raw, axis="material") if raw is not None else None
+
+
+def load_sec_type_topn_map(ckpt: dict) -> TopNMap | None:
+    """`ckpt["sec_type_topn_map"]` as a `giant.data.loader.TopNMap`, or
+    `None` if this checkpoint's `stage2_model.particle_type.target` was never
+    `"onehot"` (see `giant.pipeline.run_setup_stage`).
+
+    Pre-gitea-#29 checkpoints have no `sec_type_topn_map` key at all — before
+    #29, the secondary-species decode map and the conditioning PDG onehot map
+    were always numerically the same map, saved once under `pdg_topn_map`.
+    For those, fall back to `load_pdg_topn_map` to reproduce that exact
+    behavior; a current checkpoint always has the key (possibly `null`, if
+    `particle_type.target != "onehot"`), so this fallback never fires for one."""
+    if "sec_type_topn_map" in ckpt:
+        raw = ckpt["sec_type_topn_map"]
+        return topnmap_from_json(raw, axis="pdg") if raw is not None else None
+    return load_pdg_topn_map(ckpt)
+
+
+@dataclass(frozen=True)
+class InferenceContext:
+    """Everything needed to run a trained checkpoint forward, resolved once."""
+
+    stage1: nn.Module | None
+    stage2: nn.Module | None
+    cond_norm: Normalizer
+    tgt_norm: Normalizer
+    sec_phys_norm: Normalizer
+    pdg_map: dict[int, int]
+    mat_map: dict[str, int]
+    pdg_topn_map: TopNMap | None
+    mat_topn_map: TopNMap | None
+    sec_type_topn_map: TopNMap | None
+    particle_conditioning: str
+    material_conditioning: str
+    k_max: int
+    stage1_ddpm_steps: int
+    stage2_ddpm_steps: int
+    other_policy: str
+    model_config: dict
+    epoch: int | None
+    best_val_loss: float | None
+
+
+def load_for_inference(
+    checkpoint: Path,
+    device: torch.device,
+    command_name: str,
+    weights: str = "raw",
+    require_stage2: bool = True,
+) -> InferenceContext:
+    """Load *checkpoint* and reconstruct everything `predict`/`rollout` need
+    to run it forward, on *device*, in `eval()` mode.
+
+    *command_name* (e.g. `"predict"`/`"rollout"`) only feeds the "needs both"
+    error message below. *weights* is `"raw"` (the live training weights) or
+    `"ema"` (the EMA shadow copy, see `--ema-decay`). *require_stage2*
+    controls whether a checkpoint with an inactive stage 2
+    (`stage2_model.active = false`) is an error (both current callers need
+    both stages) or an acceptable `stage2 = None` result — kept as a real
+    parameter since `stage{1,2}_model.active` is a real, if currently
+    stage1+stage2-only-in-practice, config option.
+    """
+    ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
+    for key in ("model_config", "sec_decoder"):
+        if key not in ckpt:
+            raise CheckpointCompatibilityError(f"checkpoint has no {key} — retrain with the current code")
+
+    if "sec_phys" not in ckpt.get("normalizer", {}):
+        raise CheckpointCompatibilityError("checkpoint has no normalizer.sec_phys — retrain with the current code")
+
+    gconfig.warn_if_checkpoint_config_mismatch(checkpoint)
+
+    model_cfg = ckpt["model_config"]
+    particle_conditioning, material_conditioning = conditioning_axes(model_cfg)
+    pdg_topn_map = load_pdg_topn_map(ckpt)
+    mat_topn_map = load_mat_topn_map(ckpt)
+    if particle_conditioning == "onehot" and pdg_topn_map is None:
+        raise CheckpointCompatibilityError(
+            "checkpoint's conditioning.particle.type='onehot' but has no pdg_topn_map — retrain with the current code"
+        )
+    if material_conditioning == "onehot" and mat_topn_map is None:
+        raise CheckpointCompatibilityError(
+            "checkpoint's conditioning.material.type='onehot' but has no mat_topn_map — retrain with the current code"
+        )
+    sec_type_topn_map = load_sec_type_topn_map(ckpt)
+    particle_type_target = stage_cfg(model_cfg, "stage2").get("particle_type", {}).get("target", "onehot")
+    if particle_type_target == "onehot" and sec_type_topn_map is None:
+        raise CheckpointCompatibilityError(
+            "checkpoint's stage2_model.particle_type.target='onehot' but has no "
+            "sec_type_topn_map — retrain with the current code"
+        )
+    other_policy = particle_type_other_policy(model_cfg)
+    stage1_ddpm_steps = ddpm_steps(model_cfg, "stage1")
+    stage2_ddpm_steps = ddpm_steps(model_cfg, "stage2")
+    k_max = stage_cfg(model_cfg, "stage2").get("k_max", K_MAX)
+    pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
+    mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
+    cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
+    tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
+    sec_phys_norm = Normalizer.from_dict(ckpt["normalizer"]["sec_phys"])
+
+    built = build_models(model_cfg)
+    stage1, stage2 = built["stage1"], built["stage2"]
+    if require_stage2 and (stage1 is None or stage2 is None):
+        raise CheckpointCompatibilityError(
+            f"checkpoint has an inactive stage1 or stage2 — {command_name} needs both (see stage{{1,2}}_model.active)"
+        )
+
+    if weights == "raw":
+        model_key, sec_key = "model", "sec_decoder"
+    else:
+        model_key, sec_key = "model_ema", "sec_decoder_ema"
+        if model_key not in ckpt or sec_key not in ckpt:
+            raise CheckpointCompatibilityError(
+                f"{checkpoint} has no EMA weights (trained before --ema-decay, "
+                "or with --ema-decay 0) — use --weights raw"
+            )
+    if stage1 is not None:
+        stage1.load_state_dict(ckpt[model_key])
+        stage1.to(device).eval()
+    if stage2 is not None:
+        stage2.load_state_dict(ckpt[sec_key])
+        stage2.to(device).eval()
+
+    return InferenceContext(
+        stage1=stage1,
+        stage2=stage2,
+        cond_norm=cond_norm,
+        tgt_norm=tgt_norm,
+        sec_phys_norm=sec_phys_norm,
+        pdg_map=pdg_map,
+        mat_map=mat_map,
+        pdg_topn_map=pdg_topn_map,
+        mat_topn_map=mat_topn_map,
+        sec_type_topn_map=sec_type_topn_map,
+        particle_conditioning=particle_conditioning,
+        material_conditioning=material_conditioning,
+        k_max=k_max,
+        stage1_ddpm_steps=stage1_ddpm_steps,
+        stage2_ddpm_steps=stage2_ddpm_steps,
+        other_policy=other_policy,
+        model_config=model_cfg,
+        epoch=ckpt.get("epoch"),
+        best_val_loss=ckpt.get("best_val_loss"),
+    )

giant/cli.py

@@ -19,7 +19,6 @@ from tqdm import tqdm
 
 from giant import config as gconfig
 from giant.constants import (
-    K_MAX,
     LOCAL_TARGET_NAMES,
     PREDICT_COORD_METADATA_KEY,
     PREDICT_SCHEMA_VERSION,
@@ -39,11 +38,9 @@ from giant.data.transforms import (
     inv_local_frame_rotation,
     inv_log_transform,
     reconstruct_post_pos,
-    Normalizer,
 )
-from giant.data.setup_cache import topnmap_from_json
+from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference
 from giant.geometry import GeometryOracle
-from giant.model.network import build_models
 from giant.pipeline import run_train_job
 from giant.rollout import (
     L1DistCollector,
@@ -71,58 +68,6 @@ def _router_total_experts(router_cfg: dict) -> int:
     return int(router_cfg.get("n_experts", 1))
 
 
-def _conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]:
-    """(particle_conditioning, material_conditioning) for
-    `giant.data.transforms.build_cond_features`/`build_features` — from
-    either a v0.2 checkpoint's flat `model_config["conditioning"]` (one
-    shared string, same for both axes) or a new-format one (independent
-    `model_config["conditioning"]["particle"/"material"]["type"]` — the two
-    axes are configured independently and may differ)."""
-    raw = model_cfg.get("conditioning", default)
-    if isinstance(raw, dict):
-        return (
-            raw.get("particle", {}).get("type", default),
-            raw.get("material", {}).get("type", default),
-        )
-    return raw, raw
-
-
-def _stage_cfg(model_cfg: dict, stage: str) -> dict:
-    """`model_cfg[f"{stage}_model"]` for a new-format model_config, `{}` for
-    a v0.2 flat one (whose ddpm schedule always used `CosineSchedule`'s own
-    default `T=1000` — never a config key — and which never had
-    `particle_type` at all, so `{}` is the correct fallback for both
-    `_ddpm_steps`/`_particle_type_other_policy` below)."""
-    val = model_cfg.get(f"{stage}_model")
-    return val if isinstance(val, dict) else {}
-
-
-def _ddpm_steps(model_cfg: dict, stage: str) -> int:
-    return _stage_cfg(model_cfg, stage).get("ddpm", {}).get("n_steps", 1000)
-
-
-def _particle_type_other_policy(model_cfg: dict) -> str:
-    return _stage_cfg(model_cfg, "stage2").get("particle_type", {}).get("other_policy", "sample")
-
-
-def _load_pdg_topn_map(ckpt: dict):
-    """`ckpt["pdg_topn_map"]` as a `giant.data.loader.TopNMap`, or `None` if
-    this checkpoint's conditioning/particle_type never needed one (see
-    `giant.pipeline.run_setup_stage`, which only populates it when
-    `conditioning.particle.type` or `stage2_model.particle_type.target` is
-    `"onehot"`)."""
-    raw = ckpt.get("pdg_topn_map")
-    return topnmap_from_json(raw, axis="pdg") if raw is not None else None
-
-
-def _load_mat_topn_map(ckpt: dict):
-    """`ckpt["mat_topn_map"]` as a `giant.data.loader.TopNMap`, or `None` if
-    this checkpoint's `conditioning.material.type` was never `"onehot"` (see
-    `giant.pipeline.run_setup_stage`)."""
-    raw = ckpt.get("mat_topn_map")
-    return topnmap_from_json(raw, axis="material") if raw is not None else None
-
-
 def _batch_size_estimate_dims(model_cfg: dict, training: bool, stage: str = "stage1") -> tuple[int, int]:
     """Pick the (hidden_dim, n_blocks) that dominate per-call activation memory.
 
@@ -283,7 +228,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
 
@@ -298,34 +243,6 @@ class Weights(str, Enum):
     ema = "ema"
 
 
-def _load_model_weights(
-    model: torch.nn.Module,
-    sec_decoder: torch.nn.Module,
-    ckpt: dict,
-    weights: "Weights",
-    checkpoint_path: Path,
-) -> None:
-    """Load either the raw or EMA state dicts from a training checkpoint.
-
-    EMA weights (giant.training's shadow copy, see --ema-decay) only exist in
-    checkpoints written after that feature landed, so `ema` fails loudly
-    rather than silently falling back to raw weights a caller didn't ask for.
-    """
-    if weights == Weights.raw:
-        model_key, sec_key = "model", "sec_decoder"
-    else:
-        model_key, sec_key = "model_ema", "sec_decoder_ema"
-        if model_key not in ckpt or sec_key not in ckpt:
-            typer.echo(
-                f"error: {checkpoint_path} has no EMA weights (trained before "
-                "--ema-decay, or with --ema-decay 0) — use --weights raw",
-                err=True,
-            )
-            raise typer.Exit(1)
-    model.load_state_dict(ckpt[model_key])
-    sec_decoder.load_state_dict(ckpt[sec_key])
-
-
 @app.command()
 def train(
     data: Annotated[Path, typer.Argument(help="Parquet file or directory of parquet files")],
@@ -542,6 +459,34 @@ 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,
     val_fraction: Annotated[Optional[float], typer.Option("--val-fraction", "-f")] = None,
     seed: Annotated[
         Optional[int],
@@ -653,141 +598,61 @@ def train(
                 )
                 raise typer.Exit(1)
 
-    cli_train = {
-        k: v
-        for k, v in {
-            "epochs": epochs,
-            "batch_size": batch_size_value,
-            "lr": lr,
-            "weight_decay": weight_decay,
-            "ema_decay": ema_decay,
-            "warmup_epochs": warmup_epochs,
-            "val_fraction": val_fraction,
-            "num_workers": num_workers,
-            "seed": seed,
-            "validate_every": validate_every,
-            "validate_steps": validate_steps,
-            "max_val_batches": max_val_batches,
-            "wandb": wandb,
-            "wandb_project": wandb_project,
-            "wandb_run_name": wandb_run_name,
-            "wandb_log_every": wandb_log_every,
-        }.items()
-        if v is not None
-    }
-    # --hidden-dim/--n-blocks/--dropout are stage-1-only shorthands kept for
-    # backward compatibility (they predate stage2_model having its own
-    # flags); --stage1-*/--stage2-* below are the explicit, discoverable
-    # per-stage flags, and take precedence when both are given.
-    cli_stage1_model: dict[str, object] = {
-        k: v
-        for k, v in {
-            "hidden_dim": hidden_dim,
-            "n_res_blocks": n_blocks,
-            "dropout": dropout,
-        }.items()
-        if v is not None
-    }
-    cli_stage1_model.update(
-        {
-            k: v
-            for k, v in {
-                "hidden_dim": stage1_hidden_dim,
-                "n_res_blocks": stage1_n_res_blocks,
-                "dropout": stage1_dropout,
-            }.items()
-            if v is not None
-        }
-    )
-    cli_stage2_model: dict[str, object] = {
-        k: v
-        for k, v in {
-            "hidden_dim": stage2_hidden_dim,
-            "n_res_blocks": stage2_n_res_blocks,
-            "dropout": stage2_dropout,
-            "decoder": stage2_decoder.value if stage2_decoder is not None else None,
-            "k_max": stage2_k_max,
-            "context_dim": stage2_context_dim,
-            "stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
-        }.items()
-        if v is not None
-    }
     cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
-    if cli_router:
-        cli_stage1_model["router"] = cli_router
-
-    # --emb-dim/--conditioning set both conditioning axes (v0.2 had one
-    # shared value for particle+material).
-    cli_conditioning: dict[str, dict] = {}
-    if emb_dim is not None:
-        cli_conditioning["particle"] = {"emb_dim": emb_dim}
-        cli_conditioning["material"] = {"emb_dim": emb_dim}
-    if conditioning is not None:
-        cli_conditioning.setdefault("particle", {})["type"] = conditioning.value
-        cli_conditioning.setdefault("material", {})["type"] = conditioning.value
-
-    overrides: dict[str, dict] = {}
-    if cli_train:
-        overrides["train"] = cli_train
-    if cli_stage1_model:
-        overrides["stage1_model"] = cli_stage1_model
-    if cli_stage2_model:
-        overrides["stage2_model"] = cli_stage2_model
-    if cli_conditioning:
-        overrides["conditioning"] = cli_conditioning
-
-    # --mode/--n-critic/--gp-weight/--critic-lr/--noise-dim apply to BOTH
-    # stages by default (v0.2 had one global mode/wgan config shared by both
-    # — see giant.config.migrate_config's train.mode /
-    # train.{n_critic,gp_weight,critic_lr} precedent); the --stage1-*/
-    # --stage2-* variants below override a single stage independently (decision
-    # 7), which is what actually enables e.g. `--stage1-generator flow
-    # --stage2-generator wgan`.
-    if mode is not None:
-        overrides.setdefault("stage1_model", {})["generator"] = mode.value
-        overrides.setdefault("stage2_model", {})["generator"] = mode.value
-    if stage1_generator is not None:
-        overrides.setdefault("stage1_model", {})["generator"] = stage1_generator.value
-    if stage2_generator is not None:
-        overrides.setdefault("stage2_model", {})["generator"] = stage2_generator.value
-
-    shared_wgan_overrides = {
-        k: v
-        for k, v in {
-            "n_critic": n_critic,
-            "gp_weight": gp_weight,
-            "noise_dim": noise_dim,
-            "critic_lr": critic_lr,
-        }.items()
-        if v is not None
+    flag_values: dict[str, object] = {
+        "epochs": epochs,
+        "batch_size": batch_size_value,
+        "lr": lr,
+        "weight_decay": weight_decay,
+        "ema_decay": ema_decay,
+        "warmup_epochs": warmup_epochs,
+        "val_fraction": val_fraction,
+        "num_workers": num_workers,
+        "seed": seed,
+        "validate_every": validate_every,
+        "validate_steps": validate_steps,
+        "max_val_batches": max_val_batches,
+        "wandb": wandb,
+        "wandb_project": wandb_project,
+        "wandb_run_name": wandb_run_name,
+        "wandb_log_every": wandb_log_every,
+        "hidden_dim": hidden_dim,
+        "n_blocks": n_blocks,
+        "dropout": dropout,
+        "stage1_hidden_dim": stage1_hidden_dim,
+        "stage1_n_res_blocks": stage1_n_res_blocks,
+        "stage1_dropout": stage1_dropout,
+        "stage2_hidden_dim": stage2_hidden_dim,
+        "stage2_n_res_blocks": stage2_n_res_blocks,
+        "stage2_dropout": stage2_dropout,
+        "stage2_decoder": stage2_decoder.value if stage2_decoder is not None else None,
+        "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,
+        "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,
+        "conditioning": conditioning.value if conditioning is not None else None,
+        "emb_dim": emb_dim,
+        "router_config": cli_router or None,
+        "n_critic": n_critic,
+        "gp_weight": gp_weight,
+        "noise_dim": noise_dim,
+        "critic_lr": critic_lr,
+        "stage1_n_critic": stage1_n_critic,
+        "stage1_gp_weight": stage1_gp_weight,
+        "stage1_noise_dim": stage1_noise_dim,
+        "stage1_critic_lr": stage1_critic_lr,
+        "stage2_n_critic": stage2_n_critic,
+        "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_wgan_overrides = {
-        k: v
-        for k, v in {
-            "n_critic": stage1_n_critic,
-            "gp_weight": stage1_gp_weight,
-            "noise_dim": stage1_noise_dim,
-            "critic_lr": stage1_critic_lr,
-        }.items()
-        if v is not None
-    }
-    stage2_wgan_overrides = {
-        k: v
-        for k, v in {
-            "n_critic": stage2_n_critic,
-            "gp_weight": stage2_gp_weight,
-            "noise_dim": stage2_noise_dim,
-            "critic_lr": stage2_critic_lr,
-        }.items()
-        if v is not None
-    }
-    for stage_name, stage_specific in (
-        ("stage1_model", stage1_wgan_overrides),
-        ("stage2_model", stage2_wgan_overrides),
-    ):
-        stage_wgan = {**shared_wgan_overrides, **stage_specific}
-        if stage_wgan:
-            overrides.setdefault(stage_name, {}).setdefault("wgan", {}).update(stage_wgan)
+    overrides = gconfig.overrides_from_flags(flag_values)
 
     cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides)
     gconfig.validate_config(cfg)
@@ -912,75 +777,32 @@ def new_run(
     (with the full dataset-derived meta section), so this scaffold's meta
     section is just a placeholder recording what was asked for and when.
     """
-    cli_train = {
-        k: v
-        for k, v in {
-            "epochs": epochs,
-            "batch_size": batch_size,
-            "lr": lr,
-        }.items()
-        if v is not None
-    }
-    cli_stage1_model: dict[str, object] = {
-        k: v
-        for k, v in {
-            "hidden_dim": hidden_dim,
-            "n_res_blocks": n_blocks,
-            "dropout": dropout,
-        }.items()
-        if v is not None
-    }
-    cli_stage1_model.update(
-        {
-            k: v
-            for k, v in {
-                "hidden_dim": stage1_hidden_dim,
-                "n_res_blocks": stage1_n_res_blocks,
-                "dropout": stage1_dropout,
-            }.items()
-            if v is not None
-        }
-    )
-    cli_stage2_model: dict[str, object] = {
-        k: v
-        for k, v in {
-            "hidden_dim": stage2_hidden_dim,
-            "n_res_blocks": stage2_n_res_blocks,
-            "dropout": stage2_dropout,
-            "decoder": stage2_decoder.value if stage2_decoder is not None else None,
-            "k_max": stage2_k_max,
-            "context_dim": stage2_context_dim,
-            "stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
-        }.items()
-        if v is not None
-    }
     cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
-    if cli_router:
-        cli_stage1_model["router"] = cli_router
-    cli_conditioning: dict[str, dict] = {}
-    if emb_dim is not None:
-        cli_conditioning["particle"] = {"emb_dim": emb_dim}
-        cli_conditioning["material"] = {"emb_dim": emb_dim}
-    if conditioning is not None:
-        cli_conditioning.setdefault("particle", {})["type"] = conditioning.value
-        cli_conditioning.setdefault("material", {})["type"] = conditioning.value
-
-    overrides: dict[str, dict] = {}
-    if cli_train:
-        overrides["train"] = cli_train
-    if cli_stage1_model:
-        overrides["stage1_model"] = cli_stage1_model
-    if cli_stage2_model:
-        overrides["stage2_model"] = cli_stage2_model
-    if cli_conditioning:
-        overrides["conditioning"] = cli_conditioning
-    if mode is not None:
-        overrides.setdefault("stage1_model", {})["generator"] = mode.value
-        overrides.setdefault("stage2_model", {})["generator"] = mode.value
-    if stage1_generator is not None:
-        overrides.setdefault("stage1_model", {})["generator"] = stage1_generator.value
-    if stage2_generator is not None:
-        overrides.setdefault("stage2_model", {})["generator"] = stage2_generator.value
+    flag_values: dict[str, object] = {
+        "epochs": epochs,
+        "batch_size": batch_size,
+        "lr": lr,
+        "hidden_dim": hidden_dim,
+        "n_blocks": n_blocks,
+        "dropout": dropout,
+        "stage1_hidden_dim": stage1_hidden_dim,
+        "stage1_n_res_blocks": stage1_n_res_blocks,
+        "stage1_dropout": stage1_dropout,
+        "stage2_hidden_dim": stage2_hidden_dim,
+        "stage2_n_res_blocks": stage2_n_res_blocks,
+        "stage2_dropout": stage2_dropout,
+        "stage2_decoder": stage2_decoder.value if stage2_decoder is not None else None,
+        "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,
+        "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,
+        "conditioning": conditioning.value if conditioning is not None else None,
+        "emb_dim": emb_dim,
+        "router_config": cli_router or None,
+    }
+    overrides = gconfig.overrides_from_flags(flag_values)
 
     cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides)
     gconfig.validate_config(cfg)
@@ -1119,37 +941,26 @@ def predict(
     typer.echo(f"device: {_device}")
 
     # --- Load checkpoint ---
-    ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
-    if "model_config" not in ckpt:
-        typer.echo(
-            "error: checkpoint has no model_config — retrain with the current code",
-            err=True,
-        )
+    try:
+        ctx = load_for_inference(checkpoint, _device, "predict", weights=weights.value)
+    except CheckpointCompatibilityError as exc:
+        typer.echo(f"error: {exc}", err=True)
         raise typer.Exit(1)
+    typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})")
 
-    if "sec_decoder" not in ckpt:
-        typer.echo(
-            "error: checkpoint has no sec_decoder — retrain with the current code",
-            err=True,
-        )
-        raise typer.Exit(1)
-
-    if "sec_phys" not in ckpt.get("normalizer", {}):
-        typer.echo(
-            "error: checkpoint has no normalizer.sec_phys — retrain with the current code",
-            err=True,
-        )
-        raise typer.Exit(1)
-
-    model_cfg = ckpt["model_config"]
-    pdg_topn_map = _load_pdg_topn_map(ckpt)
-    mat_topn_map = _load_mat_topn_map(ckpt)
-    other_policy = _particle_type_other_policy(model_cfg)
-    stage1_ddpm_steps = _ddpm_steps(model_cfg, "stage1")
-    stage2_k_max = _stage_cfg(model_cfg, "stage2").get("k_max", K_MAX)
+    assert ctx.stage1 is not None and ctx.stage2 is not None  # require_stage2=True (default) guarantees this
+    model, sec_decoder = ctx.stage1, ctx.stage2
+    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
+    stage2_k_max = ctx.k_max
 
     if batch_size_auto:
-        est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(model_cfg, training=False)
+        est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(ctx.model_config, training=False)
         try:
             batch_size_value = gconfig.estimate_batch_size(
                 est_hidden_dim,
@@ -1164,42 +975,6 @@ def predict(
 
     assert batch_size_value is not None
     bs = batch_size_value
-    particle_conditioning, material_conditioning = _conditioning_axes(model_cfg)
-    if particle_conditioning == "onehot" and pdg_topn_map is None:
-        typer.echo(
-            "error: checkpoint's conditioning.particle.type='onehot' but has "
-            "no pdg_topn_map — retrain with the current code",
-            err=True,
-        )
-        raise typer.Exit(1)
-    if material_conditioning == "onehot" and mat_topn_map is None:
-        typer.echo(
-            "error: checkpoint's conditioning.material.type='onehot' but has "
-            "no mat_topn_map — retrain with the current code",
-            err=True,
-        )
-        raise typer.Exit(1)
-    pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
-    mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
-    cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
-    tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
-    sec_phys_norm = Normalizer.from_dict(ckpt["normalizer"]["sec_phys"])
-
-    built = build_models(model_cfg)
-    model, sec_decoder = built["stage1"], built["stage2"]
-    if model is None or sec_decoder is None:
-        typer.echo(
-            "error: checkpoint has an inactive stage1 or stage2 — giant "
-            "predict needs both (see stage{1,2}_model.active)",
-            err=True,
-        )
-        raise typer.Exit(1)
-    _load_model_weights(model, sec_decoder, ckpt, weights, checkpoint)
-    model.to(_device).eval()
-    sec_decoder.to(_device).eval()
-
-    typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})")
-    gconfig.warn_if_checkpoint_config_mismatch(checkpoint)
 
     # --- Output path ---
     out, dataset_path, pred_uuid = _resolve_prediction_output(data, out)
@@ -1221,8 +996,12 @@ def predict(
             return iter_file_chunks(path, offset=offset, k_max=stage2_k_max)
         return iter_cond_chunks(path, offset=offset)
 
-    cond_pdg_topn = pdg_topn_map.class_map if particle_conditioning == "onehot" else None
-    cond_mat_topn = mat_topn_map.class_map if material_conditioning == "onehot" else None
+    # load_for_inference already guarantees pdg_topn_map/mat_topn_map are not
+    # None whenever the matching conditioning axis is "onehot" — the extra
+    # `is not None` conjuncts below are redundant at runtime, just narrowing
+    # for the type checker.
+    cond_pdg_topn = pdg_topn_map.class_map if pdg_topn_map is not None and particle_conditioning == "onehot" else None
+    cond_mat_topn = mat_topn_map.class_map if mat_topn_map is not None and material_conditioning == "onehot" else None
 
     def _concat(a: dict[str, np.ndarray], b: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
         return {k: np.concatenate([a[k], b[k]], axis=0) for k in a}
@@ -1231,7 +1010,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,
@@ -1241,7 +1020,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,
@@ -1325,7 +1106,7 @@ def predict(
                 piece["pre_dir"],
                 sec_phys_norm,
                 pdg_map,
-                pdg_topn_map,
+                sec_type_topn_map,
                 other_policy,
                 None,
             )
@@ -1532,65 +1313,25 @@ def rollout(
     _device = torch.device(device) if device else gconfig.auto_device()
     typer.echo(f"device: {_device}")
 
-    ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
-    for key in ("model_config", "sec_decoder"):
-        if key not in ckpt:
-            typer.echo(
-                f"error: checkpoint has no {key} — retrain with the current code",
-                err=True,
-            )
-            raise typer.Exit(1)
-
-    if "sec_phys" not in ckpt.get("normalizer", {}):
-        typer.echo(
-            "error: checkpoint has no normalizer.sec_phys — retrain with the current code",
-            err=True,
-        )
+    try:
+        ctx = load_for_inference(checkpoint, _device, "rollout", weights=weights.value)
+    except CheckpointCompatibilityError as exc:
+        typer.echo(f"error: {exc}", err=True)
         raise typer.Exit(1)
+    typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})")
 
-    gconfig.warn_if_checkpoint_config_mismatch(checkpoint)
     training_cfg = gconfig.load_checkpoint_config(checkpoint)
 
-    model_cfg = ckpt["model_config"]
-    particle_conditioning, material_conditioning = _conditioning_axes(model_cfg)
-    pdg_topn_map = _load_pdg_topn_map(ckpt)
-    mat_topn_map = _load_mat_topn_map(ckpt)
-    if particle_conditioning == "onehot" and pdg_topn_map is None:
-        typer.echo(
-            "error: checkpoint's conditioning.particle.type='onehot' but has "
-            "no pdg_topn_map — retrain with the current code",
-            err=True,
-        )
-        raise typer.Exit(1)
-    if material_conditioning == "onehot" and mat_topn_map is None:
-        typer.echo(
-            "error: checkpoint's conditioning.material.type='onehot' but has "
-            "no mat_topn_map — retrain with the current code",
-            err=True,
-        )
-        raise typer.Exit(1)
-    other_policy = _particle_type_other_policy(model_cfg)
-    stage1_ddpm_steps = _ddpm_steps(model_cfg, "stage1")
-    stage2_ddpm_steps = _ddpm_steps(model_cfg, "stage2")
-    pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
-    mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
-    cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
-    tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
-    sec_phys_norm = Normalizer.from_dict(ckpt["normalizer"]["sec_phys"])
-
-    built = build_models(model_cfg)
-    model, sec_decoder = built["stage1"], built["stage2"]
-    if model is None or sec_decoder is None:
-        typer.echo(
-            "error: checkpoint has an inactive stage1 or stage2 — giant "
-            "rollout needs both (see stage{1,2}_model.active)",
-            err=True,
-        )
-        raise typer.Exit(1)
-    _load_model_weights(model, sec_decoder, ckpt, weights, checkpoint)
-    model.to(_device).eval()
-    sec_decoder.to(_device).eval()
-    typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})")
+    assert ctx.stage1 is not None and ctx.stage2 is not None  # require_stage2=True (default) guarantees this
+    model, sec_decoder = ctx.stage1, ctx.stage2
+    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
+    model_cfg = ctx.model_config
 
     oracle = GeometryOracle.load(geometry)
     typer.echo(f"loaded geometry oracle: {geometry} (escape_threshold={oracle.escape_threshold:.3f})")
@@ -1647,6 +1388,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,
@@ -1687,8 +1429,8 @@ 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),
-            "training_epoch": ckpt.get("epoch"),
-            "best_val_loss": ckpt.get("best_val_loss"),
+            "training_epoch": ctx.epoch,
+            "best_val_loss": ctx.best_val_loss,
             # [train]/[meta] from the sibling config.toml (giant.config.save_config)
             # — empty dicts if the checkpoint has no config.toml next to it.
             "training_config": dict(training_cfg.get("train", {})),

giant/config.py

@@ -14,10 +14,12 @@ from pathlib import Path
 import numpy as np
 import torch
 
+from giant._migration import V02_FIXED_FACTS, V02_MODEL_KEY_TO_STAGES, reject_legacy_router_expert_sizing
+
 
 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 +49,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)
@@ -384,26 +384,24 @@ class NSecConfig:
     # 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")
+    # 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"
 
     @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"),
         )
 
     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}
 
 
 @dataclass(frozen=True)
@@ -420,6 +418,11 @@ 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
 
     @classmethod
     def from_dict(cls, d: dict | None) -> "ParticleTypeConfig":
@@ -428,10 +431,16 @@ 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),
         )
 
     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,
+        }
 
 
 @dataclass(frozen=True)
@@ -893,6 +902,121 @@ def _deep_merge(base: dict, override: dict) -> dict:
     return result
 
 
+@dataclass(frozen=True)
+class FlagSpec:
+    """One CLI flag's mapping into the config-overrides tree.
+
+    `paths` lists every dotted config path this flag writes (>1 means fan-out
+    to multiple stages/axes, e.g. `--mode` -> both stages' `generator`).
+    `precedence` controls write order when two flags target the same path:
+    specs are applied in ascending precedence, so a higher-precedence (more
+    specific) flag overwrites a lower-precedence (shared/shorthand) one —
+    this is the "build a shared dict, then let a more specific dict win"
+    pattern `giant train`/`giant new-run` need (e.g. `--hidden-dim` vs
+    `--stage1-hidden-dim`, or `--n-critic` vs `--stage1-n-critic`),
+    generalized to one mechanism instead of three different ad hoc ones.
+    """
+
+    name: str
+    paths: tuple[str, ...]
+    precedence: int = 0
+
+
+# Flag -> config-path table shared by `giant train`/`giant new-run`
+# (giant/cli.py) so both commands resolve CLI overrides identically. See
+# issues.md Issue 3: this replaces ~140 lines of hand-written, imperative
+# dict-building in cli.py with one declarative table plus
+# `overrides_from_flags` below.
+FLAG_SPECS: tuple[FlagSpec, ...] = (
+    # train block -- flat pass-through, unique paths, precedence irrelevant.
+    FlagSpec("epochs", ("train.epochs",)),
+    FlagSpec("batch_size", ("train.batch_size",)),
+    FlagSpec("lr", ("train.lr",)),
+    FlagSpec("weight_decay", ("train.weight_decay",)),
+    FlagSpec("ema_decay", ("train.ema_decay",)),
+    FlagSpec("warmup_epochs", ("train.warmup_epochs",)),
+    FlagSpec("val_fraction", ("train.val_fraction",)),
+    FlagSpec("num_workers", ("train.num_workers",)),
+    FlagSpec("seed", ("train.seed",)),
+    FlagSpec("validate_every", ("train.validate_every",)),
+    FlagSpec("validate_steps", ("train.validate_steps",)),
+    FlagSpec("max_val_batches", ("train.max_val_batches",)),
+    FlagSpec("wandb", ("train.wandb",)),
+    FlagSpec("wandb_project", ("train.wandb_project",)),
+    FlagSpec("wandb_run_name", ("train.wandb_run_name",)),
+    FlagSpec("wandb_log_every", ("train.wandb_log_every",)),
+    # --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.
+    FlagSpec("hidden_dim", ("stage1_model.hidden_dim",), precedence=0),
+    FlagSpec("stage1_hidden_dim", ("stage1_model.hidden_dim",), precedence=1),
+    FlagSpec("n_blocks", ("stage1_model.n_res_blocks",), precedence=0),
+    FlagSpec("stage1_n_res_blocks", ("stage1_model.n_res_blocks",), precedence=1),
+    FlagSpec("dropout", ("stage1_model.dropout",), precedence=0),
+    FlagSpec("stage1_dropout", ("stage1_model.dropout",), precedence=1),
+    # stage2-only knobs.
+    FlagSpec("stage2_hidden_dim", ("stage2_model.hidden_dim",)),
+    FlagSpec("stage2_n_res_blocks", ("stage2_model.n_res_blocks",)),
+    FlagSpec("stage2_dropout", ("stage2_model.dropout",)),
+    FlagSpec("stage2_decoder", ("stage2_model.decoder",)),
+    FlagSpec("stage2_k_max", ("stage2_model.k_max",)),
+    FlagSpec("stage2_context_dim", ("stage2_model.context_dim",)),
+    FlagSpec("stage2_stage1_context", ("stage2_model.stage1_context",)),
+    # --mode applies to both stages by default (v0.2 had one shared
+    # mode/wgan config); --stage{1,2}-generator override a single stage.
+    FlagSpec("mode", ("stage1_model.generator", "stage2_model.generator"), precedence=0),
+    FlagSpec("stage1_generator", ("stage1_model.generator",), precedence=1),
+    FlagSpec("stage2_generator", ("stage2_model.generator",), precedence=1),
+    # --emb-dim/--conditioning set both conditioning axes (v0.2 had one
+    # shared value for particle+material).
+    FlagSpec("conditioning", ("conditioning.particle.type", "conditioning.material.type")),
+    FlagSpec("emb_dim", ("conditioning.particle.emb_dim", "conditioning.material.emb_dim")),
+    # Pre-aggregated router override dict (built by `_router_cli_overrides`
+    # in cli.py from --router/--router-type/--n-experts/--router-axis).
+    # Router overrides only ever land on stage1_model -- this asymmetry is
+    # deliberate (see cli.py) and must not be "fixed" into a fan-out here.
+    FlagSpec("router_config", ("stage1_model.router",)),
+    # WGAN: shared knobs apply to both stages by default (v0.2 had one
+    # shared wgan config); --stage{1,2}-* override a single stage.
+    FlagSpec("n_critic", ("stage1_model.wgan.n_critic", "stage2_model.wgan.n_critic"), precedence=0),
+    FlagSpec("stage1_n_critic", ("stage1_model.wgan.n_critic",), precedence=1),
+    FlagSpec("stage2_n_critic", ("stage2_model.wgan.n_critic",), precedence=1),
+    FlagSpec("gp_weight", ("stage1_model.wgan.gp_weight", "stage2_model.wgan.gp_weight"), precedence=0),
+    FlagSpec("stage1_gp_weight", ("stage1_model.wgan.gp_weight",), precedence=1),
+    FlagSpec("stage2_gp_weight", ("stage2_model.wgan.gp_weight",), precedence=1),
+    FlagSpec("noise_dim", ("stage1_model.wgan.noise_dim", "stage2_model.wgan.noise_dim"), precedence=0),
+    FlagSpec("stage1_noise_dim", ("stage1_model.wgan.noise_dim",), precedence=1),
+    FlagSpec("stage2_noise_dim", ("stage2_model.wgan.noise_dim",), precedence=1),
+    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",)),
+)
+
+
+def overrides_from_flags(values: dict[str, object]) -> dict:
+    """Build the nested, section-keyed config-overrides dict
+    `merge_cli_overrides` expects, from `{flag_name: value}`.
+
+    Flags absent from `values`, or mapped to `None` (= not given on the
+    CLI), are skipped. See `FlagSpec`/`FLAG_SPECS` above for the precedence
+    rule applied when two flags target the same path.
+    """
+    overrides: dict = {}
+    for spec in sorted(FLAG_SPECS, key=lambda s: s.precedence):
+        if spec.name not in values or values[spec.name] is None:
+            continue
+        for path in spec.paths:
+            _set_path(overrides, path, values[spec.name])
+    return overrides
+
+
 # v0.2 [train] keys that pass through to v0.3 [train] unchanged (same name,
 # same meaning) when present in the loaded file — everything model-shaped
 # moved to the stage/conditioning blocks instead (see the rest of
@@ -916,14 +1040,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 = (
@@ -983,7 +1099,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])
@@ -1000,18 +1116,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
@@ -1019,21 +1124,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
@@ -1181,6 +1274,14 @@ def validate_config(cfg: dict) -> None:
             "(standalone stage-2 evaluation only, never for rollout)"
         )
 
+    if _get_path(cfg, "stage2_model.stage1_context") == "sampled":
+        raise ValueError(
+            "stage2_model.stage1_context = 'sampled' is accepted by the schema "
+            "but not implemented — trainers.py always trains stage 2 against "
+            "the ground-truth stage-1 output; use 'truth' (default) instead "
+            "(see issues.md Issue 16 for the planned implementation)"
+        )
+
     if (
         _get_path(cfg, "stage2_model.n_sec.mode") == "truth"
         and _get_path(cfg, "stage1_model.active")
@@ -1197,6 +1298,13 @@ 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)"
+            )
         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'")

giant/data/dataset.py

@@ -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
@@ -129,17 +144,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 +160,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:
                     (
@@ -226,14 +231,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:

giant/data/loader.py

@@ -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

giant/data/transforms.py

@@ -1,4 +1,5 @@
 import warnings
+from typing import NamedTuple
 
 import numpy as np
 
@@ -860,36 +861,10 @@ def _cond_normalizer_transform(
     return ((cond_cont - mean) / std).astype(np.float32)
 
 
-def build_features(
-    data: dict[str, np.ndarray],
-    pdg_map: dict[int, int],
-    mat_map: dict[str, int],
-    cond_normalizer: Normalizer | None = None,
-    target_normalizer: Normalizer | None = None,
-    sec_phys_normalizer: Normalizer | None = None,
-    fit: bool = False,
-    proc_map: dict[str, int] | None = None,
-    require_secondaries: bool = False,
-    particle_conditioning: str = "embedding",
-    material_conditioning: str = "embedding",
-    sec_phys_only: bool = False,
-    pdg_topn_map: dict[int, int] | None = None,
-    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.
+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)
@@ -908,6 +883,40 @@ def build_features(
               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],
+    mat_map: dict[str, int],
+    cond_normalizer: Normalizer | None = None,
+    target_normalizer: Normalizer | None = None,
+    sec_phys_normalizer: Normalizer | None = None,
+    fit: bool = False,
+    proc_map: dict[str, int] | None = None,
+    require_secondaries: bool = False,
+    particle_conditioning: str = "embedding",
+    material_conditioning: str = "embedding",
+    sec_phys_only: bool = False,
+    pdg_topn_map: dict[int, int] | None = None,
+    mat_topn_map: dict[str, int] | None = None,
+    sec_type_class_map: dict | None = None,
+    k_max: int = K_MAX,
+) -> 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
@@ -1054,14 +1063,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,
     )

giant/model/_legacy.py

@@ -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

giant/model/builders.py

@@ -0,0 +1,211 @@
+"""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.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["conditioning"]
+    particle_cfg = conditioning["particle"]
+    material_cfg = conditioning["material"]
+    particle_conditioning = particle_cfg["type"]
+    conditioning_cfg = ConditioningConfig.from_dict(conditioning)
+    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
+        # 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 generator != "wgan" 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,
+            n_sec_head_k_max=n_sec_head_k_max,
+            cond_enc=shared_cond_enc,
+        )
+
+    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
+        # wgan has no time_dim concept — see the matching comment in stage 1
+        # above.
+        time_dim = getattr(s2_spec, generator).time_dim if generator != "wgan" else 64
+        n_sec_owner = s2_spec.n_sec.owner
+        k_max = s2_spec.k_max
+        particle_type_cfg = s2_spec.particle_type.to_dict()
+
+        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,
+                build_n_sec_head=n_sec_owner != "stage1",
+                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,
+            )
+        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,
+                build_n_sec_head=n_sec_owner != "stage1",
+                particle_type_cfg=particle_type_cfg,
+                cond_enc=shared_cond_enc,
+            )
+
+    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["conditioning"]
+    particle_cfg = conditioning["particle"]
+    material_cfg = conditioning["material"]
+    conditioning_cfg = ConditioningConfig.from_dict(conditioning)
+    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 s1_spec.generator == "wgan":
+        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",
+        )
+
+    if s2_spec.active and s2_spec.generator == "wgan":
+        k_max = s2_spec.k_max
+        particle_type_cfg = s2_spec.particle_type.to_dict()
+        in_dim = stage2_trunk_sec_dim(
+            particle_type_cfg, "wgan", 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,
+        )
+
+    return result

giant/model/encoders.py

@@ -0,0 +1,122 @@
+"""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.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM
+from giant.model.layers import _make_axis_mlp
+
+
+def cat_col_layout(particle_type: str, material_type: str) -> tuple[int | None, int | None]:
+    """`cond_cat` column indices for each axis's top-N-onehot index, or
+    `None` if that axis isn't `"onehot"`.
+
+    Columns 0/1 are always the dense pdg/material vocab index. The particle
+    top-N column (if any) comes next, then the material top-N column (if
+    any) — `giant.data.transforms.build_cond_features`/`build_features`
+    append columns in this same order, so the two sides must never drift
+    apart.
+    """
+    col = 2
+    particle_col = None
+    if particle_type == "onehot":
+        particle_col = col
+        col += 1
+    material_col = None
+    if material_type == "onehot":
+        material_col = col
+        col += 1
+    return particle_col, material_col
+
+
+class ConditionEncoder(nn.Module):
+    """Fuses continuous conditioning with particle/material identity.
+
+    The particle and material axes are configured independently
+    (`particle_cfg`/`material_cfg`, each `{"type", "emb_dim", "n_layers"}`)
+    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[:, COND_DIM_BASE:]` — 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) — see `_cat_col_layout`.
+    """
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        particle_cfg: dict,
+        material_cfg: dict,
+        cont_dim: int = COND_DIM,
+        out_dim: int = 128,
+    ) -> None:
+        super().__init__()
+        self.particle_cfg = dict(particle_cfg)
+        self.material_cfg = dict(material_cfg)
+        self._particle_topn_col, self._material_topn_col = cat_col_layout(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.get("n_layers", 1))
+        elif p_type != "onehot":
+            raise ValueError(f"unknown conditioning.particle.type {p_type!r}")
+
+        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.get("n_layers", 1))
+        elif m_type != "onehot":
+            raise ValueError(f"unknown conditioning.material.type {m_type!r}")
+
+        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[:, 0])
+        if p_type == "physical":
+            particle_phys = cond_cont[:, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM]
+            return self.particle_mlp(particle_phys)
+        assert self._particle_topn_col is not None
+        return F.one_hot(
+            cond_cat[:, self._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[:, 1])
+        if m_type == "physical":
+            material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :]
+            return self.material_mlp(material_phys)
+        assert self._material_topn_col is not None
+        return F.one_hot(
+            cond_cat[:, self._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[:, :COND_DIM_BASE], pdg_e, mat_e], dim=-1)
+        return self.mlp(x)

giant/model/history.py

@@ -0,0 +1,154 @@
+"""History encoders — stage-2 autoregressive only. Self-contained, no
+dependency on any other `giant.model` submodule (issues.md Issue 8)."""
+
+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 implementations. 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); encoders that need
+    incremental state for that path additionally implement `init_cache`/
+    `step` (see `AttentionHistory`) — `MarkovHistory` doesn't need to, since
+    its per-step cost is already O(1) (it only ever looks at the previous
+    token, not the full prefix)."""
+
+    def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
+
+
+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
+
+
+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,
+        token_feat: torch.Tensor,
+        has_prev: torch.Tensor,
+        cache: list[torch.Tensor | None],
+    ) -> tuple[torch.Tensor, list[torch.Tensor | None]]:
+        """`token_feat`/`has_prev`: `(B, 1, in_dim)`/`(B, 1)` — the newest
+        token's own features (what would be `feat[:, k]` in `forward`).
+        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."""
+        x = self._embed(token_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

giant/model/layers.py

@@ -0,0 +1,76 @@
+"""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)
+
+
+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))
+
+
+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

giant/model/models.py

@@ -0,0 +1,593 @@
+"""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.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 AttentionHistory, HistoryEncoder, MarkovHistory
+from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding
+from giant.model.routers import Router
+from giant.model.trunks import build_trunk
+
+# ---------------------------------------------------------------------------
+# Stage models
+# ---------------------------------------------------------------------------
+
+
+def resolve_type_n_classes(particle_type_cfg: dict, 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.get("target", "physical") == "onehot":
+        return particle_type_cfg.get("n_classes", 0) or particle_emb_dim
+    return particle_emb_dim
+
+
+def stage2_type_dim(particle_type_cfg: dict, 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)."""
+    target = particle_type_cfg.get("target", "physical")
+    return PARTICLE_PHYS_DIM if target == "physical" else emb_dim
+
+
+def stage2_trunk_sec_dim(particle_type_cfg: dict, 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 `generator == "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)`. Under
+    `generator in ("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`.
+    """
+    target = particle_type_cfg.get("target", "physical")
+    if target == "physical":
+        return k_max * SEC_SLOT_DIM
+    if generator == "wgan":
+        return k_max * (CONT_SLOT_DIM + emb_dim)
+    return k_max * CONT_SLOT_DIM
+
+
+class Stage1Model(nn.Module):
+    """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`).
+
+    `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: dict,
+        material_cfg: dict,
+        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,
+        n_sec_head_k_max: int | None = None,
+        cond_enc: ConditionEncoder | None = None,
+    ) -> None:
+        super().__init__()
+        self.generator_kind = generator
+        self.noise_dim = noise_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)
+        )
+        has_time = generator in ("flow", "ddpm")
+        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 = noise_dim if generator == "wgan" else x_dim
+        self.trunk = build_trunk(router, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout)
+        self.n_sec_head = None
+        if n_sec_head_k_max is not None:
+            self.n_sec_head = nn.Sequential(
+                nn.Linear(cond_out_dim, hidden_dim // 2),
+                nn.SiLU(),
+                nn.Linear(hidden_dim // 2, n_sec_head_k_max + 1),
+            )
+
+    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 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."""
+        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')"
+            )
+        c_emb = self.cond_enc(cond_cont, cond_cat)
+        return self.n_sec_head(c_emb)
+
+
+class Stage2OneShot(nn.Module):
+    """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 `generator in ("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 `generator == "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.
+
+    `cond_enc`, if given, is used in place of building a fresh
+    `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`).
+    """
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        particle_cfg: dict,
+        material_cfg: dict,
+        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,
+        build_n_sec_head: bool = True,
+        particle_type_cfg: dict | None = None,
+        cond_enc: ConditionEncoder | None = None,
+    ) -> None:
+        super().__init__()
+        self.generator_kind = generator
+        self.noise_dim = noise_dim
+        self.k_max = k_max
+        self.particle_type_cfg = dict(particle_type_cfg or {"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)
+        )
+        self.context_adapter = ContextAdapter(x_dim, context_dim)
+        self.fuse = nn.Sequential(
+            nn.Linear(cond_out_dim + context_dim, cond_out_dim),
+            nn.SiLU(),
+        )
+        has_time = generator in ("flow", "ddpm")
+        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 = noise_dim if generator == "wgan" else sec_dim
+        self.trunk = build_trunk(router, in_dim, sec_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout)
+        self.n_sec_head = None
+        if build_n_sec_head:
+            self.n_sec_head = nn.Sequential(
+                nn.Linear(cond_out_dim, hidden_dim // 2),
+                nn.SiLU(),
+                nn.Linear(hidden_dim // 2, k_max + 1),
+            )
+        self.type_head = None
+        target = self.particle_type_cfg.get("target", "physical")
+        if target != "physical" and generator in ("flow", "ddpm"):
+            emb_dim = resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"])
+            self.type_head = nn.Sequential(
+                nn.Linear(cond_out_dim, hidden_dim // 2),
+                nn.SiLU(),
+                nn.Linear(hidden_dim // 2, k_max * emb_dim),
+            )
+            self._type_k_max = k_max
+            self._type_emb_dim = emb_dim
+
+    def _cond_embed(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.fuse(torch.cat([base, ctx], dim=-1))
+
+    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:
+        if self.n_sec_head is None:
+            raise RuntimeError(
+                "this Stage2OneShot 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"
+            )
+        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)."""
+        if self.type_head is None:
+            raise RuntimeError(
+                "this Stage2OneShot 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)"
+            )
+        c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
+        return self.type_head(c_emb).view(-1, self._type_k_max, self._type_emb_dim)
+
+
+class Stage2Autoregressive(nn.Module):
+    """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`/the trunk.
+
+    `cond_enc`, if given, is used in place of building a fresh
+    `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`).
+    """
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        particle_cfg: dict,
+        material_cfg: dict,
+        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,
+        build_n_sec_head: bool = True,
+        particle_type_cfg: dict | None = None,
+        history: str = "markov",
+        attn_n_heads: int = 4,
+        attn_n_layers: int = 2,
+        cond_enc: ConditionEncoder | None = None,
+    ) -> None:
+        super().__init__()
+        if history not in ("markov", "attention"):
+            raise ValueError(f"stage2_model.autoregressive.history={history!r} — must be 'markov' or 'attention'")
+        self.history_kind = history
+        self.generator_kind = generator
+        self.noise_dim = noise_dim
+        self.k_max = k_max
+        self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"})
+        emb_dim = resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"])
+        self.type_dim = stage2_type_dim(self.particle_type_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)
+        )
+        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 = (
+            AttentionHistory(hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers)
+            if history == "attention"
+            else MarkovHistory(hist_in_dim, history_dim)
+        )
+        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(),
+        )
+
+        has_time = generator in ("flow", "ddpm")
+        self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
+        merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
+        token_dim = stage2_trunk_sec_dim(self.particle_type_cfg, generator, 1, emb_dim)
+        in_dim = noise_dim if generator == "wgan" else token_dim
+        self.trunk = build_trunk(
+            router,
+            in_dim,
+            token_dim,
+            hidden_dim,
+            n_res_blocks,
+            merged_cond_dim,
+            dropout,
+        )
+
+        self.n_sec_head = None
+        if build_n_sec_head:
+            self.n_sec_head = nn.Sequential(
+                nn.Linear(cond_out_dim, hidden_dim // 2),
+                nn.SiLU(),
+                nn.Linear(hidden_dim // 2, k_max + 1),
+            )
+        self.type_head = None
+        target = self.particle_type_cfg.get("target", "physical")
+        if target != "physical" and generator in ("flow", "ddpm"):
+            self.type_head = nn.Sequential(
+                nn.Linear(cond_out_dim, hidden_dim // 2),
+                nn.SiLU(),
+                nn.Linear(hidden_dim // 2, self.type_dim),
+            )
+
+    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): `None` under `history="markov"` (its
+        per-step cost is already O(1) — see `HistoryEncoder`'s docstring), or
+        `AttentionHistory.init_cache()` under `history="attention"`."""
+        if isinstance(self.history_encoder, AttentionHistory):
+            return self.history_encoder.init_cache()
+        return None
+
+    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."""
+        if isinstance(self.history_encoder, AttentionHistory):
+            return self.history_encoder.step(token_feat, has_prev, cache)
+        return self.history_encoder(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:
+        if self.n_sec_head is None:
+            raise RuntimeError(
+                "this Stage2Autoregressive 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"
+            )
+        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:
+        if self.type_head is None:
+            raise RuntimeError(
+                "this Stage2Autoregressive 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)"
+            )
+        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)
+
+
+class CriticModel(nn.Module):
+    """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`). Used only when that stage's `generator == "wgan"`."""
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        particle_cfg: dict,
+        material_cfg: dict,
+        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,
+    ) -> None:
+        super().__init__()
+        if stage not in ("stage1", "stage2"):
+            raise ValueError(f"stage must be 'stage1' or 'stage2', got {stage!r}")
+        self.stage = stage
+        self.cond_enc = ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
+        if stage == "stage2":
+            self.context_adapter = ContextAdapter(context_in_dim, context_dim)
+            self.fuse = nn.Sequential(
+                nn.Linear(cond_out_dim + context_dim, cond_out_dim),
+                nn.SiLU(),
+            )
+        self.input_proj = nn.Linear(in_dim, hidden_dim)
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_res_blocks)])
+        self.out_norm = nn.LayerNorm(hidden_dim)
+        self.out_proj = nn.Linear(hidden_dim, 1)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        base = self.cond_enc(cond_cont, cond_cat)
+        if self.stage == "stage2":
+            ctx = self.context_adapter(stage1_out)
+            cond = self.fuse(torch.cat([base, ctx], dim=-1))
+        else:
+            cond = base
+        h = self.input_proj(x)
+        for block in self.blocks:
+            h = block(h, cond)
+        return self.out_proj(self.out_norm(h)).squeeze(-1)

giant/model/routers.py

@@ -0,0 +1,375 @@
+"""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.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.
+        """
+        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."""
+        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)."""
+        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."""
+        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("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[:, 0])  # (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[:, 0])
+        mat_e = self.mat_emb(cond_cat[:, 1])
+        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

giant/model/trunks.py

@@ -0,0 +1,155 @@
+"""Trunks: everything downstream of the fused conditioning vector — monolithic
+or expert-routed (issues.md Issue 8)."""
+
+import torch
+import torch.nn as nn
+
+from giant.model.layers import ResBlock
+from giant.model.routers import Router
+
+
+class ExpertTrunk(nn.Module):
+    """One small expert: `input_proj -> ResBlock stack -> out_proj`.
+
+    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,
+    ) -> None:
+        super().__init__()
+        self.input_proj = nn.Linear(in_dim, hidden_dim)
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_blocks)])
+        self.out_proj = nn.Linear(hidden_dim, out_dim)
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        x = self.input_proj(x)
+        for block in self.blocks:
+            x = block(x, cond)
+        return self.out_proj(x)
+
+
+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.
+    """
+    if training:
+        weights = router.combine_weights(cond_cont, cond_cat)  # (B, n_experts)
+        out = torch.zeros(x.shape[0], experts[0].out_proj.out_features, device=x.device)
+        for i, expert in enumerate(experts):
+            out = out + weights[:, i : i + 1] * expert(x, cond)
+        return out
+
+    idx = router.top1(cond_cont, cond_cat)  # (B,)
+    out_dim = experts[0].out_proj.out_features
+    out = torch.zeros(x.shape[0], out_dim, device=x.device)
+    for i, expert in enumerate(experts):
+        mask = idx == i
+        if mask.any():
+            out[mask] = expert(x[mask], cond[mask])
+    return out
+
+
+class Trunk(nn.Module):
+    """Interface implemented by `MonolithicTrunk`/`RoutedTrunk`: everything
+    downstream of the fused conditioning vector, i.e. the actual generative
+    trunk of a stage (`input_proj -> blocks -> out_proj`, monolithic or
+    expert-routed)."""
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+
+class MonolithicTrunk(Trunk):
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        hidden_dim: int,
+        n_res_blocks: int,
+        cond_dim: int,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.input_proj = nn.Linear(in_dim, hidden_dim)
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_res_blocks)])
+        self.out_proj = nn.Linear(hidden_dim, out_dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+    ) -> torch.Tensor:
+        x = self.input_proj(x)
+        for block in self.blocks:
+            x = block(x, cond)
+        return self.out_proj(x)
+
+
+class RoutedTrunk(Trunk):
+    def __init__(
+        self,
+        router: Router,
+        in_dim: int,
+        out_dim: int,
+        hidden_dim: int,
+        n_res_blocks: int,
+        cond_dim: int,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.router = router
+        self.experts = nn.ModuleList(
+            [ExpertTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) 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,
+    in_dim: int,
+    out_dim: int,
+    hidden_dim: int,
+    n_res_blocks: int,
+    cond_dim: int,
+    dropout: float = 0.0,
+) -> Trunk:
+    if router is not None:
+        return RoutedTrunk(router, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
+    return MonolithicTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)

giant/particles.py

@@ -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"`

giant/pipeline.py

@@ -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_dict = cfg["stage2_model"].get("particle_type") or {}
+    particle_type_target = config.ParticleTypeConfig.from_dict(particle_type_cfg_dict).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_dict, 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,
@@ -259,6 +273,10 @@ def run_setup_stage(
                     sec_phys_only=True,
                     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 +310,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,
@@ -381,8 +400,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:
@@ -486,6 +505,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,

giant/rollout.py

@@ -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]:
@@ -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.get("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,
@@ -690,7 +701,7 @@ def _step_chunk(
         tr["pre_dir"],
         sec_phys_norm,
         pdg_map,
-        pdg_topn_map,
+        sec_type_topn_map,
         other_policy,
         rng,
     )

giant/tools/bump_dataset_version.py

@@ -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)
 # ---------------------------------------------------------------------------
 
 

giant/tools/check_migration_v02_v03.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").

giant/tools/create_root_files.py

@@ -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$")
 

giant/tools/dwarf.py

@@ -1,6 +1,6 @@
 """dwarf — little helper to `giant`: dataset/tooling CLI for the geant_steps pipeline.
 
-Unifies the standalone scripts/*.py conversion, migration, versioning, and
+Unifies the standalone giant/tools/*.py conversion, migration, versioning, and
 simulation-fanout tools into one Typer app so there's a single command name
 (and `--help`) to remember instead of five differently-hyphenated ones.
 """
@@ -14,20 +14,20 @@ import typer
 from typing_extensions import Annotated
 
 from giant.config import Conditioning
-from scripts.bump_dataset_version import (
+from giant.tools.bump_dataset_version import (
     run_bump_gen,
     run_bump_schema,
     run_create_manifest,
     run_status,
     run_update_manifest,
 )
-from scripts.create_root_files import run_make_root
-from scripts.geometry_oracle import run_build_geometry_oracle
-from scripts.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
-from scripts.migrate_geant_steps import run_migration
-from scripts.steps_to_parquet import convert_steps_to_parquet
-from scripts.steps_to_parquet_parallel import run_parallel_job
-from scripts.warm_setup_cache import run_warm_setup_cache
+from giant.tools.create_root_files import run_make_root
+from giant.tools.geometry_oracle import run_build_geometry_oracle
+from giant.tools.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
+from giant.tools.migrate_geant_steps import run_migration
+from giant.tools.steps_to_parquet import convert_steps_to_parquet
+from giant.tools.steps_to_parquet_parallel import run_parallel_job
+from giant.tools.warm_setup_cache import run_warm_setup_cache
 
 app = typer.Typer(no_args_is_help=True)
 

giant/tools/profile_analysis_costs.py

@@ -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

giant/tools/steps_to_parquet_parallel.py

@@ -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

giant/training/loop.py

@@ -109,6 +109,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,
@@ -142,6 +143,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,
     }

giant/training/trainers.py

@@ -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,7 +24,8 @@ 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.data.dataset import StepBatch
+from giant.model.network import Router, resolve_type_n_classes, stage2_type_dim
 from giant.model.schedule import (
     CosineSchedule,
     flow_matching_loss,
@@ -72,8 +74,8 @@ 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))
 
 
 @dataclass(frozen=True)
@@ -96,7 +98,7 @@ class StageSpec:
 
     # 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
 
     # optimization
     lr: float = 3e-4
@@ -149,7 +151,9 @@ class StageSpec:
             lambda_weight=stage_spec.lambda_weight,
             n_sec_lambda=s2_spec.n_sec.lambda_weight,
             particle_type=s2_spec.particle_type,
-            particle_type_emb_dim=cfg["conditioning"]["particle"]["emb_dim"],
+            particle_type_n_classes=resolve_type_n_classes(
+                s2_spec.particle_type.to_dict(), cfg["conditioning"]["particle"]["emb_dim"]
+            ),
             # 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
@@ -185,11 +189,10 @@ 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
+    Reads only the shared `StepBatch` (`giant.data.dataset`) — 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.
@@ -237,7 +240,7 @@ class StageTrainer:
         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
+        self.particle_type_n_classes = spec.particle_type_n_classes
         self.ema_decay = spec.ema_decay
 
         self.ema_model: torch.nn.Module | None = None
@@ -255,10 +258,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 ------------------------------------------------
@@ -328,7 +331,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 +357,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
 
@@ -575,7 +578,7 @@ 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
@@ -615,9 +618,12 @@ 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
         if self.spec.lambda_balance > 0:
@@ -639,7 +645,7 @@ class FlowDDPMStageTrainer(StageTrainer):
             "nsec_acc": nsec_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,
@@ -659,7 +665,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 +680,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).
 
@@ -737,7 +753,7 @@ 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):
+    def _stage2_real_and_fake(self, batch_tensors: _Stage2RealFakeBatch, 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
@@ -745,7 +761,7 @@ class WGANStageTrainer(StageTrainer):
         multiplied on the fake side yet."""
         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)
 
@@ -777,7 +793,7 @@ class WGANStageTrainer(StageTrainer):
 
         return real, fake_raw, mask, critic_fn
 
-    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,
@@ -802,7 +818,7 @@ class WGANStageTrainer(StageTrainer):
             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),
+                _Stage2RealFakeBatch(cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx),
                 stage1_ctx,
                 global_step,
                 device,
@@ -824,7 +840,7 @@ class WGANStageTrainer(StageTrainer):
                     fake_raw,
                     sec_cont.size(1),
                     CONT_SLOT_DIM,
-                    stage2_type_dim(self.particle_type_cfg, self.particle_type_emb_dim),
+                    stage2_type_dim(self.particle_type_cfg, self.particle_type_n_classes),
                     tau,
                     grad_probe=grad_probe,
                 )

giant/validate.py

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

pyproject.toml

@@ -50,13 +50,13 @@ analysis = [
 
 [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 +70,7 @@ requires = ["hatchling"]
 build-backend = "hatchling.build"
 
 [tool.hatch.build.targets.wheel]
-packages = ["giant", "scripts"]
+packages = ["giant"]
 
 [tool.uv]
 conflicts = [

tests/test_bump_dataset_version.py

@@ -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

tests/test_checkpoint_io.py

@@ -0,0 +1,280 @@
+"""Tests for giant.checkpoint_io.load_for_inference (issues.md Issue 5) —
+the shared bootstrap `giant predict`/`giant rollout` use to go from a
+checkpoint path to ready-to-run models."""
+
+from __future__ import annotations
+
+import copy
+
+import numpy as np
+import pytest
+import torch
+
+from giant import config as gconfig
+from giant.checkpoint_io import (
+    CheckpointCompatibilityError,
+    InferenceContext,
+    conditioning_axes,
+    load_for_inference,
+    stage_cfg,
+)
+from giant.data.loader import TopNMap
+from giant.data.setup_cache import topnmap_to_json
+from giant.data.transforms import Normalizer
+from giant.model.network import build_models
+
+PDG_MAP = {11: 0, 22: 1, -11: 2}
+MAT_MAP = {"G4_PbWO4": 0, "G4_AIR": 1}
+
+
+def _model_cfg(stage2_active: bool = True) -> dict:
+    """DEFAULT_CONFIG-derived, shrunk for speed — same pattern as
+    tests/test_network.py::_minimal_model_config. Default `conditioning`
+    (both axes "physical") needs no top-N vocab map, so this is a cheap,
+    fully self-contained happy-path config."""
+    cfg = copy.deepcopy(gconfig.DEFAULT_CONFIG)
+    cfg["conditioning"]["particle"]["emb_dim"] = 4
+    cfg["conditioning"]["material"]["emb_dim"] = 4
+    cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1})
+    cfg["stage2_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "k_max": 3})
+    cfg["stage2_model"]["active"] = stage2_active
+    return {
+        "pdg_vocab": len(PDG_MAP),
+        "mat_vocab": len(MAT_MAP),
+        "conditioning": cfg["conditioning"],
+        "stage1_model": cfg["stage1_model"],
+        "stage2_model": cfg["stage2_model"],
+    }
+
+
+def _norms() -> tuple[Normalizer, Normalizer, Normalizer]:
+    rng = np.random.default_rng(0)
+    cond = Normalizer().fit(rng.standard_normal((100, 15)).astype(np.float32))
+    tgt = Normalizer().fit(rng.standard_normal((100, 9)).astype(np.float32))
+    sec_phys = Normalizer().fit(rng.standard_normal((100, 2)).astype(np.float32))
+    return cond, tgt, sec_phys
+
+
+def _write_checkpoint(tmp_path, model_cfg=None, ema: bool = False, **ckpt_overrides):
+    cfg = model_cfg if model_cfg is not None else _model_cfg()
+    built = build_models(cfg)
+    stage1, stage2 = built["stage1"], built["stage2"]
+    cond, tgt, sec_phys = _norms()
+    ckpt: dict = {
+        "model_config": cfg,
+        "model": stage1.state_dict() if stage1 is not None else {},
+        "sec_decoder": stage2.state_dict() if stage2 is not None else {},
+        "pdg_map": PDG_MAP,
+        "mat_map": MAT_MAP,
+        "normalizer": {"cond": cond.to_dict(), "target": tgt.to_dict(), "sec_phys": sec_phys.to_dict()},
+        "epoch": 3,
+        "best_val_loss": 0.5,
+    }
+    if ema:
+        ckpt["model_ema"] = stage1.state_dict() if stage1 is not None else {}
+        ckpt["sec_decoder_ema"] = stage2.state_dict() if stage2 is not None else {}
+    # DEFAULT_CONFIG's stage2_model.particle_type.target defaults to
+    # "onehot", and giant train's pipeline (gitea #29) now always writes a
+    # sec_type_topn_map in that case — default one in here too, unless a
+    # test explicitly overrides it, so fixtures represent a real, loadable
+    # checkpoint by default rather than exercising the "missing" guard by
+    # accident.
+    particle_type_target = cfg.get("stage2_model", {}).get("particle_type", {}).get("target", "onehot")
+    if particle_type_target == "onehot" and "sec_type_topn_map" not in ckpt_overrides:
+        default_sec_type_topn = TopNMap(class_map=dict(zip(PDG_MAP, range(len(PDG_MAP)))), other_members={})
+        ckpt["sec_type_topn_map"] = topnmap_to_json(default_sec_type_topn)
+    ckpt.update(ckpt_overrides)
+    path = tmp_path / "ckpt.pt"
+    torch.save(ckpt, path)
+    return path
+
+
+def _onehot_model_cfg() -> dict:
+    cfg = _model_cfg()
+    cfg["conditioning"]["particle"]["type"] = "onehot"
+    return cfg
+
+
+# ---------------------------------------------------------------------------
+# Happy path
+# ---------------------------------------------------------------------------
+
+
+def test_happy_path_returns_populated_context(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+    assert isinstance(ctx, InferenceContext)
+    assert ctx.stage1 is not None and ctx.stage2 is not None
+    assert not ctx.stage1.training
+    assert not ctx.stage2.training
+    assert next(ctx.stage1.parameters()).device == torch.device("cpu")
+    assert ctx.pdg_map == PDG_MAP
+    assert ctx.mat_map == MAT_MAP
+    assert all(isinstance(k, int) for k in ctx.pdg_map)
+    assert all(isinstance(k, str) for k in ctx.mat_map)
+    assert ctx.particle_conditioning == "physical"
+    assert ctx.material_conditioning == "physical"
+    assert ctx.k_max == 3
+    assert ctx.epoch == 3
+    assert ctx.best_val_loss == 0.5
+    assert ctx.model_config["stage1_model"]["hidden_dim"] == 8
+
+
+def test_happy_path_normalizer_values_round_trip(tmp_path):
+    cond, tgt, sec_phys = _norms()
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+    assert ctx.cond_norm.mean is not None and cond.mean is not None
+    assert ctx.tgt_norm.mean is not None and tgt.mean is not None
+    assert ctx.sec_phys_norm.mean is not None and sec_phys.mean is not None
+    np.testing.assert_allclose(ctx.cond_norm.mean, cond.mean)
+    np.testing.assert_allclose(ctx.tgt_norm.mean, tgt.mean)
+    np.testing.assert_allclose(ctx.sec_phys_norm.mean, sec_phys.mean)
+
+
+# ---------------------------------------------------------------------------
+# Guards
+# ---------------------------------------------------------------------------
+
+
+def test_missing_model_config_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["model_config"]
+    torch.save(ckpt, checkpoint)
+
+    with pytest.raises(CheckpointCompatibilityError, match="no model_config"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_missing_sec_decoder_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["sec_decoder"]
+    torch.save(ckpt, checkpoint)
+
+    with pytest.raises(CheckpointCompatibilityError, match="no sec_decoder"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_missing_sec_phys_normalizer_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["normalizer"]["sec_phys"]
+    torch.save(ckpt, checkpoint)
+
+    with pytest.raises(CheckpointCompatibilityError, match="no normalizer.sec_phys"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_onehot_particle_conditioning_without_topn_map_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, model_cfg=_onehot_model_cfg())
+
+    with pytest.raises(CheckpointCompatibilityError, match="pdg_topn_map"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_onehot_particle_conditioning_with_topn_map_succeeds(tmp_path):
+    topn = TopNMap(class_map={11: 0, 22: 1}, other_members={})
+    checkpoint = _write_checkpoint(
+        tmp_path,
+        model_cfg=_onehot_model_cfg(),
+        pdg_topn_map=topnmap_to_json(topn),
+    )
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+    assert ctx.particle_conditioning == "onehot"
+    assert ctx.pdg_topn_map is not None
+    assert ctx.pdg_topn_map.class_map == {11: 0, 22: 1}
+
+
+def test_onehot_particle_type_target_without_sec_type_topn_map_raises(tmp_path):
+    """DEFAULT_CONFIG's stage2_model.particle_type.target="onehot" needs a
+    sec_type_topn_map (gitea #29) — a checkpoint with neither key at all
+    (not even the pre-#29 pdg_topn_map to fall back to) must fail loudly."""
+    checkpoint = _write_checkpoint(tmp_path, sec_type_topn_map=None)
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["sec_type_topn_map"]
+    torch.save(ckpt, checkpoint)
+
+    with pytest.raises(CheckpointCompatibilityError, match="sec_type_topn_map"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_pre_gitea_29_checkpoint_falls_back_to_pdg_topn_map_for_sec_type(tmp_path):
+    """A checkpoint written before gitea #29 has no sec_type_topn_map key at
+    all — conditioning and secondary-type onehot maps were always the same
+    map, saved once under pdg_topn_map. load_for_inference must reproduce
+    that exact pre-#29 behavior for such a checkpoint."""
+    topn = TopNMap(class_map={11: 0, 22: 1, -11: 2}, other_members={})
+    checkpoint = _write_checkpoint(
+        tmp_path,
+        model_cfg=_onehot_model_cfg(),
+        pdg_topn_map=topnmap_to_json(topn),
+        sec_type_topn_map=None,
+    )
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["sec_type_topn_map"]
+    torch.save(ckpt, checkpoint)
+
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+    assert ctx.sec_type_topn_map is not None
+    assert ctx.sec_type_topn_map.class_map == {11: 0, 22: 1, -11: 2}
+
+
+def test_ema_weights_requested_but_missing_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, ema=False)
+
+    with pytest.raises(CheckpointCompatibilityError, match="no EMA weights"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict", weights="ema")
+
+
+def test_ema_weights_requested_and_present_succeeds(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, ema=True)
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict", weights="ema")
+    assert ctx.stage1 is not None and ctx.stage2 is not None
+
+
+@pytest.mark.parametrize("command_name", ["predict", "rollout"])
+def test_inactive_stage_with_require_stage2_raises_with_command_name(tmp_path, command_name):
+    checkpoint = _write_checkpoint(tmp_path, model_cfg=_model_cfg(stage2_active=False))
+
+    with pytest.raises(CheckpointCompatibilityError, match=f"{command_name} needs both"):
+        load_for_inference(checkpoint, torch.device("cpu"), command_name)
+
+
+def test_inactive_stage_with_require_stage2_false_succeeds_with_stage2_none(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, model_cfg=_model_cfg(stage2_active=False))
+
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict", require_stage2=False)
+    assert ctx.stage1 is not None
+    assert ctx.stage2 is None
+
+
+# ---------------------------------------------------------------------------
+# conditioning_axes / stage_cfg
+# ---------------------------------------------------------------------------
+
+
+def test_conditioning_axes_v02_flat_string_applies_to_both_axes():
+    assert conditioning_axes({"conditioning": "embedding"}) == ("embedding", "embedding")
+
+
+def test_conditioning_axes_v03_nested_dict_independent_per_axis():
+    model_cfg = {"conditioning": {"particle": {"type": "onehot"}, "material": {"type": "physical"}}}
+    assert conditioning_axes(model_cfg) == ("onehot", "physical")
+
+
+def test_conditioning_axes_missing_key_uses_default():
+    assert conditioning_axes({}, default="embedding") == ("embedding", "embedding")
+
+
+def test_stage_cfg_new_shape_returns_subdict():
+    model_cfg = {"stage2_model": {"k_max": 7}}
+    assert stage_cfg(model_cfg, "stage2") == {"k_max": 7}
+
+
+def test_stage_cfg_v02_flat_shape_returns_empty_dict():
+    model_cfg = {"hidden_dim": 32, "n_blocks": 4}
+    assert stage_cfg(model_cfg, "stage2") == {}

tests/test_cli_predict.py

@@ -1,13 +1,18 @@
 import uuid
 
+import torch
 import yaml
+from typer.testing import CliRunner
 
 from giant.cli import (
     _CEPH_PREDICTIONS,
     _resolve_prediction_output,
     _write_prediction_ref,
+    app,
 )
 
+runner = CliRunner()
+
 
 # ---------------------------------------------------------------------------
 # _resolve_prediction_output
@@ -150,3 +155,20 @@ def test_ref_checkpoint_path_is_absolute(tmp_path):
     data = yaml.safe_load(ref_path.read_text())
 
     assert data["checkpoint"].startswith("/")
+
+
+# ---------------------------------------------------------------------------
+# Bootstrap failure surfaces via the CLI (issues.md Issue 5 — confirms
+# CheckpointCompatibilityError -> typer.Exit(1) actually wires up end-to-end,
+# not just at the giant.checkpoint_io unit level).
+# ---------------------------------------------------------------------------
+
+
+def test_predict_exits_1_on_checkpoint_missing_model_config(tmp_path):
+    checkpoint = tmp_path / "bad.pt"
+    torch.save({"sec_decoder": {}, "normalizer": {"sec_phys": {}}}, checkpoint)
+
+    result = runner.invoke(app, ["predict", "dummy.parquet", "--checkpoint", str(checkpoint)])
+
+    assert result.exit_code == 1
+    assert "checkpoint has no model_config" in result.output

tests/test_cli_rollout.py

@@ -0,0 +1,33 @@
+"""Thin CLI smoke coverage for `giant rollout` (issues.md Issue 5) — confirms
+the CheckpointCompatibilityError raised by giant.checkpoint_io.load_for_inference
+surfaces as a clean typer.Exit(1) with the expected message, end-to-end
+through the CLI, not just at the giant.checkpoint_io unit level."""
+
+from __future__ import annotations
+
+import torch
+from typer.testing import CliRunner
+
+from giant.cli import app
+
+runner = CliRunner()
+
+
+def test_rollout_exits_1_on_checkpoint_missing_model_config(tmp_path):
+    checkpoint = tmp_path / "bad.pt"
+    torch.save({"sec_decoder": {}, "normalizer": {"sec_phys": {}}}, checkpoint)
+
+    result = runner.invoke(
+        app,
+        [
+            "rollout",
+            "dummy.parquet",
+            "--checkpoint",
+            str(checkpoint),
+            "--geometry",
+            "dummy_geometry.pkl",
+        ],
+    )
+
+    assert result.exit_code == 1
+    assert "checkpoint has no model_config" in result.output

tests/test_cli_train_overrides.py

@@ -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
 
@@ -94,3 +95,83 @@ def test_wgan_knobs_split_per_stage(monkeypatch, tmp_path):
     assert cfg["stage1_model"]["wgan"]["gp_weight"] == 10.0
     assert cfg["stage2_model"]["wgan"]["n_critic"] == 5
     assert cfg["stage2_model"]["wgan"]["gp_weight"] == 2.5
+
+
+def test_batch_size_invalid_string_errors(monkeypatch, tmp_path):
+    monkeypatch.setattr(cli, "run_train_job", lambda *a, **kw: None)
+    result = runner.invoke(
+        cli.app,
+        ["train", "dummy.parquet", "--out", str(tmp_path / "run"), "--batch-size", "not-a-number"],
+    )
+    assert result.exit_code == 1
+    assert "--batch-size must be an integer or 'auto'" in result.output
+
+
+def test_out_dir_resolution_prefers_explicit_out_over_resume(monkeypatch, tmp_path):
+    captured: dict = {}
+
+    def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
+        captured["out_dir"] = out_dir
+
+    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+
+    resume_dir = tmp_path / "resumed_run"
+    resume_dir.mkdir()
+    (resume_dir / "last.pt").touch()
+    explicit_out = tmp_path / "explicit_run"
+
+    result = runner.invoke(
+        cli.app,
+        ["train", "dummy.parquet", "--out", str(explicit_out), "--resume", str(resume_dir / "last.pt")],
+    )
+    assert result.exit_code == 0, result.output
+    assert captured["out_dir"] == explicit_out
+
+
+def test_out_dir_resolution_falls_back_to_resume_parent(monkeypatch, tmp_path):
+    captured: dict = {}
+
+    def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
+        captured["out_dir"] = out_dir
+
+    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+
+    resume_dir = tmp_path / "resumed_run"
+    resume_dir.mkdir()
+    (resume_dir / "last.pt").touch()
+
+    result = runner.invoke(cli.app, ["train", "dummy.parquet", "--resume", str(resume_dir / "last.pt")])
+    assert result.exit_code == 0, result.output
+    assert captured["out_dir"] == resume_dir
+
+
+def test_out_dir_resolution_defaults_when_neither_out_nor_resume_given(monkeypatch, tmp_path):
+    captured: dict = {}
+
+    def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
+        captured["out_dir"] = out_dir
+
+    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.chdir(tmp_path)
+
+    result = runner.invoke(cli.app, ["train", "dummy.parquet"])
+    assert result.exit_code == 0, result.output
+    assert captured["out_dir"] == Path("checkpoints") / cli.gconfig.default_out_dir_name(cli.gconfig.DEFAULT_CONFIG)
+
+
+def test_batch_size_auto_estimates_and_echoes(monkeypatch, tmp_path):
+    captured: dict = {}
+
+    def _fake_run_train_job(*, data, cfg, out_dir, num_workers, **kwargs):
+        captured["batch_size"] = cfg["train"]["batch_size"]
+
+    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr(cli.gconfig, "estimate_batch_size", lambda hidden_dim, n_blocks, device: 123)
+
+    result = runner.invoke(
+        cli.app,
+        ["train", "dummy.parquet", "--out", str(tmp_path / "run"), "--batch-size", "auto"],
+    )
+    assert result.exit_code == 0, result.output
+    assert captured["batch_size"] == 123
+    assert "batch_size: 123 (auto-estimated from free GPU memory)" in result.output

tests/test_config.py

@@ -75,6 +75,16 @@ def test_stage2_model_config_defaults_match_documented_v030_intent():
     assert spec.particle_type.target == "onehot"
 
 
+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
+    assert spec.to_dict()["n_classes"] == 32
+
+
 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 +105,15 @@ 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"}
 
 
 # ---------------------------------------------------------------------------
@@ -682,6 +697,15 @@ def test_validate_config_stop_token_not_implemented():
         assert "stop_token" in str(e)
 
 
+def test_validate_config_stage1_context_sampled_not_implemented():
+    cfg = _cfg_with(**{"stage2_model.stage1_context": "sampled"})
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "sampled" in str(e)
+
+
 def test_validate_config_n_sec_truth_rejected_for_rollout_capable_checkpoint():
     """'n_sec.mode = "truth" is invalid for a rollout-capable checkpoint' —
     both stages active means giant rollout
@@ -733,6 +757,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 +835,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",
         }
@@ -889,6 +939,142 @@ def test_merge_cli_overrides_real_config_fixtures_pass_key_validation(fixture_na
     gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / fixture_name, {})  # must not raise
 
 
+# ---------------------------------------------------------------------------
+# overrides_from_flags (issues.md Issue 3): the flag -> config-path table
+# shared by `giant train`/`giant new-run`. Each test below pins one
+# precedence rule directly, without CliRunner — see also
+# tests/test_cli_train_overrides.py for the thin end-to-end smoke coverage.
+# ---------------------------------------------------------------------------
+
+
+def test_overrides_from_flags_empty_values_yield_empty_overrides():
+    assert gconfig.overrides_from_flags({}) == {}
+    assert gconfig.overrides_from_flags({"epochs": None, "hidden_dim": None}) == {}
+
+
+def test_overrides_from_flags_train_block_passthrough():
+    overrides = gconfig.overrides_from_flags({"epochs": 5, "lr": 1e-3, "hidden_dim": None})
+    assert overrides == {"train": {"epochs": 5, "lr": 1e-3}}
+
+
+@pytest.mark.parametrize(
+    ("shorthand", "explicit", "path_key"),
+    [
+        ("hidden_dim", "stage1_hidden_dim", "hidden_dim"),
+        ("n_blocks", "stage1_n_res_blocks", "n_res_blocks"),
+        ("dropout", "stage1_dropout", "dropout"),
+    ],
+)
+def test_overrides_from_flags_stage1_explicit_overrides_shorthand(shorthand, explicit, path_key):
+    overrides = gconfig.overrides_from_flags({shorthand: 1, explicit: 2})
+    assert overrides["stage1_model"][path_key] == 2
+
+
+@pytest.mark.parametrize(
+    ("shorthand", "path_key"),
+    [("hidden_dim", "hidden_dim"), ("n_blocks", "n_res_blocks"), ("dropout", "dropout")],
+)
+def test_overrides_from_flags_stage1_shorthand_alone(shorthand, path_key):
+    overrides = gconfig.overrides_from_flags({shorthand: 7})
+    assert overrides["stage1_model"][path_key] == 7
+
+
+def test_overrides_from_flags_stage2_only_knobs():
+    overrides = gconfig.overrides_from_flags(
+        {
+            "stage2_hidden_dim": 32,
+            "stage2_n_res_blocks": 4,
+            "stage2_dropout": 0.1,
+            "stage2_decoder": "one_shot",
+            "stage2_k_max": 8,
+            "stage2_context_dim": 16,
+            "stage2_stage1_context": "sampled",
+        }
+    )
+    assert overrides["stage2_model"] == {
+        "hidden_dim": 32,
+        "n_res_blocks": 4,
+        "dropout": 0.1,
+        "decoder": "one_shot",
+        "k_max": 8,
+        "context_dim": 16,
+        "stage1_context": "sampled",
+    }
+    assert "stage1_model" not in overrides
+
+
+def test_overrides_from_flags_mode_fans_to_both_stages():
+    overrides = gconfig.overrides_from_flags({"mode": "wgan"})
+    assert overrides["stage1_model"]["generator"] == "wgan"
+    assert overrides["stage2_model"]["generator"] == "wgan"
+
+
+def test_overrides_from_flags_stage1_generator_overrides_mode_for_stage1_only():
+    overrides = gconfig.overrides_from_flags({"mode": "wgan", "stage1_generator": "flow"})
+    assert overrides["stage1_model"]["generator"] == "flow"
+    assert overrides["stage2_model"]["generator"] == "wgan"
+
+
+def test_overrides_from_flags_stage2_generator_overrides_mode_for_stage2_only():
+    overrides = gconfig.overrides_from_flags({"mode": "wgan", "stage2_generator": "flow"})
+    assert overrides["stage1_model"]["generator"] == "wgan"
+    assert overrides["stage2_model"]["generator"] == "flow"
+
+
+def test_overrides_from_flags_emb_dim_sets_both_conditioning_axes():
+    overrides = gconfig.overrides_from_flags({"emb_dim": 24})
+    assert overrides["conditioning"]["particle"]["emb_dim"] == 24
+    assert overrides["conditioning"]["material"]["emb_dim"] == 24
+
+
+def test_overrides_from_flags_conditioning_sets_both_axes_type():
+    overrides = gconfig.overrides_from_flags({"conditioning": "onehot"})
+    assert overrides["conditioning"]["particle"]["type"] == "onehot"
+    assert overrides["conditioning"]["material"]["type"] == "onehot"
+
+
+def test_overrides_from_flags_router_config_only_touches_stage1():
+    overrides = gconfig.overrides_from_flags({"router_config": {"enabled": True, "type": "energy"}})
+    assert overrides["stage1_model"]["router"] == {"enabled": True, "type": "energy"}
+    assert "stage2_model" not in overrides
+
+
+@pytest.mark.parametrize(
+    ("shared", "stage1_specific", "stage2_specific", "path_key"),
+    [
+        ("n_critic", "stage1_n_critic", "stage2_n_critic", "n_critic"),
+        ("gp_weight", "stage1_gp_weight", "stage2_gp_weight", "gp_weight"),
+        ("noise_dim", "stage1_noise_dim", "stage2_noise_dim", "noise_dim"),
+        ("critic_lr", "stage1_critic_lr", "stage2_critic_lr", "critic_lr"),
+    ],
+)
+def test_overrides_from_flags_wgan_knobs_split_per_stage(shared, stage1_specific, stage2_specific, path_key):
+    overrides = gconfig.overrides_from_flags({shared: 5.0, stage1_specific: 3.0})
+    assert overrides["stage1_model"]["wgan"][path_key] == 3.0
+    assert overrides["stage2_model"]["wgan"][path_key] == 5.0
+
+    overrides = gconfig.overrides_from_flags({shared: 5.0, stage2_specific: 2.5})
+    assert overrides["stage1_model"]["wgan"][path_key] == 5.0
+    assert overrides["stage2_model"]["wgan"][path_key] == 2.5
+
+
+@pytest.mark.parametrize(
+    ("stage_flag", "stage_model", "path_key"),
+    [
+        ("stage1_critic_hidden_dim", "stage1_model", "critic_hidden_dim"),
+        ("stage1_critic_n_res_blocks", "stage1_model", "critic_n_res_blocks"),
+        ("stage2_critic_hidden_dim", "stage2_model", "critic_hidden_dim"),
+        ("stage2_critic_n_res_blocks", "stage2_model", "critic_n_res_blocks"),
+    ],
+)
+def test_overrides_from_flags_critic_sizing_is_stage_scoped_only(stage_flag, stage_model, path_key):
+    """critic_hidden_dim/critic_n_res_blocks are architectural per-stage
+    knobs (gitea #28) — unlike n_critic/gp_weight/noise_dim/critic_lr above,
+    there is deliberately no shared alias that fans out to both stages."""
+    overrides = gconfig.overrides_from_flags({stage_flag: 32})
+    assert overrides == {stage_model: {"wgan": {path_key: 32}}}
+
+
 # ---------------------------------------------------------------------------
 # checkpoint config-mismatch warnings (unchanged surface, still exercised)
 # ---------------------------------------------------------------------------

tests/test_config_consumed_keys.py

@@ -0,0 +1,156 @@
+"""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
+
+_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.stage1_context": (
+        "issues.md Issue 1 — trainers.py hardcodes stage1_ctx to the "
+        "ground-truth stage-1 output; 'sampled' is now rejected loudly by "
+        "validate_config (not silently accepted), but the key still isn't "
+        "read by any build/train consumer file since only 'truth' can pass "
+        "validation — see Issue 16 for the real implementation"
+    ),
+    "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 _leaf_paths(node: dict, prefix: str = "") -> list[str]:
+    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 _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."
+    )

tests/test_create_root_files.py

@@ -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

tests/test_dataset.py

@@ -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."""

tests/test_dwarf.py

@@ -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

tests/test_migration_v02_v03.py

@@ -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

tests/test_network.py

@@ -288,6 +288,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 = {"type": "physical", "emb_dim": 6, "n_layers": 1}
+    particle_type_cfg = {"target": "onehot", "lambda": 1.0, "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 -----------------------------------------------------------
 
 
@@ -423,6 +450,29 @@ 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 = {"type": "physical", "emb_dim": 6, "n_layers": 1}
+    particle_type_cfg = {"target": "onehot", "lambda": 1.0, "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
+
+
 @pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
 @pytest.mark.parametrize("generator", ["wgan", "flow"])
 @pytest.mark.parametrize("history", ["markov", "attention"])
@@ -655,6 +705,33 @@ 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_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.input_proj.in_features > default_n_classes_critic.input_proj.in_features
+
+
 # ── build_models/build_critics: DEFAULT_CONFIG fallback drift (issues.md #1) ─
 
 
@@ -708,3 +785,44 @@ def test_build_critics_omitted_particle_type_matches_default_config():
     # 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.input_proj.out_features == 8
+    assert len(inherited.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.input_proj.out_features == 16
+    assert len(overridden.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.input_proj.out_features == 8
+    assert len(inherited.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.input_proj.out_features == 16
+    assert len(overridden.blocks) == 3

tests/test_pipeline.py

@@ -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 —

tests/test_rollout.py

@@ -430,7 +430,7 @@ def _run_v3(
     max_tracks_per_event=100,
     seeds=None,
     conditioning="physical",
-    pdg_topn_map=None,
+    sec_type_topn_map=None,
     other_policy="sample",
     seed=0,
     stage1_ddpm_steps=1000,
@@ -457,7 +457,7 @@ def _run_v3(
         escape_threshold=escape_threshold,
         particle_conditioning=conditioning,
         material_conditioning=conditioning,
-        pdg_topn_map=pdg_topn_map,
+        sec_type_topn_map=sec_type_topn_map,
         other_policy=other_policy,
         seed=seed,
         stage1_ddpm_steps=stage1_ddpm_steps,
@@ -532,7 +532,7 @@ def test_rollout_onehot_target_end_to_end(fake_material_props, decoder):
     (giant.particles.particle_phys_array) become the secondary's identity —
     unlike "physical", not just a reporting label."""
     s1, s2 = _models_v3(decoder=decoder, target="onehot", emb_dim=3)
-    rec = _run_v3(s1, s2, pdg_topn_map=PDG_TOPN_MAP, other_policy="modal")
+    rec = _run_v3(s1, s2, sec_type_topn_map=PDG_TOPN_MAP, other_policy="modal")
     assert len(rec["event_id"]) > 0
     # Every spawned secondary's nominal pdg must be one decode_topn_class can
     # actually produce (the topn map's known classes + its "other" members).
@@ -543,8 +543,8 @@ def test_rollout_onehot_target_end_to_end(fake_material_props, decoder):
 
 def test_rollout_onehot_target_missing_topn_map_raises(fake_material_props):
     s1, s2 = _models_v3(target="onehot", emb_dim=3)
-    with pytest.raises(RuntimeError, match="pdg_topn_map"):
-        _run_v3(s1, s2, pdg_topn_map=None)
+    with pytest.raises(RuntimeError, match="sec_type_topn_map"):
+        _run_v3(s1, s2, sec_type_topn_map=None)
 
 
 # --- conditioning.{particle,material}.type = "onehot" — a separate axis from
@@ -627,6 +627,86 @@ def test_rollout_conditioning_onehot_material_missing_topn_map_raises(
         _run_onehot_conditioning(mat_topn_map=None)
 
 
+SEC_TYPE_TOPN_MAP_DIFFERENT_N = TopNMap(class_map={22: 0, 11: 1, -11: 2, 13: 3}, other_members={2112: 3, 2212: 1})
+
+
+def _run_conditioning_and_type_onehot_different_n_classes():
+    """Both conditioning.particle.type="onehot" and
+    stage2_model.particle_type.target="onehot" active at once, with
+    stage2_model.particle_type.n_classes deliberately different from
+    conditioning.particle.emb_dim (gitea #29)."""
+    cond_emb_dim = len(PDG_MAP)  # 3
+    type_n_classes = 5  # deliberately different from cond_emb_dim
+    particle_cfg = {"type": "onehot", "emb_dim": cond_emb_dim, "n_layers": 1}
+    material_cfg = {"type": "onehot", "emb_dim": len(MAT_MAP), "n_layers": 1}
+    particle_type_cfg = {"target": "onehot", "n_classes": type_n_classes}
+
+    s1 = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+    ).eval()
+    sec_dim = stage2_trunk_sec_dim(particle_type_cfg, "flow", K_MAX, type_n_classes)
+    s2 = Stage2OneShot(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+        sec_dim=sec_dim,
+        generator="flow",
+        time_dim=16,
+        k_max=K_MAX,
+        particle_type_cfg=particle_type_cfg,
+    ).eval()
+    # Sanity: the model's own type_dim followed n_classes, not cond_emb_dim.
+    assert s2.type_dim == type_n_classes
+
+    cond, tgt, sec_phys = _norms()
+    return rollout(
+        s1,
+        s2,
+        _oracle(),
+        _seeds(),
+        cond,
+        tgt,
+        sec_phys,
+        PDG_MAP,
+        MAT_MAP,
+        energy_cutoff=1.0,
+        max_steps=15,
+        steps=3,
+        batch_size=128,
+        max_tracks_per_event=100,
+        escape_threshold=1e9,
+        particle_conditioning="onehot",
+        material_conditioning="onehot",
+        pdg_topn_map=COND_PDG_TOPN_MAP,
+        mat_topn_map=COND_MAT_TOPN_MAP,
+        sec_type_topn_map=SEC_TYPE_TOPN_MAP_DIFFERENT_N,
+        other_policy="modal",
+    )
+
+
+def test_rollout_conditioning_and_type_onehot_with_different_n_classes(fake_material_props):
+    """gitea #29 end-to-end: conditioning.particle.type="onehot" and
+    stage2_model.particle_type.target="onehot" now use independently sized
+    top-N maps (stage2_model.particle_type.n_classes != conditioning.particle
+    .emb_dim), and rollout must decode secondaries using the type-side map,
+    not silently reuse the conditioning-side one (the pre-#29 bug)."""
+    rec = _run_conditioning_and_type_onehot_different_n_classes()
+    assert len(rec["event_id"]) > 0
+    possible = set(SEC_TYPE_TOPN_MAP_DIFFERENT_N.class_map.keys()) | set(
+        SEC_TYPE_TOPN_MAP_DIFFERENT_N.other_members.keys()
+    )
+    secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())
+    assert secondary_pdgs <= possible
+
+
 @pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"])
 def test_rollout_embedding_target_end_to_end(decoder):
     """particle_type.target="embedding" L1-snaps to the nearest row of the

tests/test_steps_to_parquet.py

@@ -1,6 +1,6 @@
 import polars as pl
 
-from scripts import steps_to_parquet
+from giant.tools import steps_to_parquet
 
 
 def _frame() -> pl.DataFrame:

tests/test_steps_to_parquet_parallel.py

@@ -2,7 +2,7 @@ import json
 import sys
 from pathlib import Path
 
-from scripts import steps_to_parquet_parallel
+from giant.tools import steps_to_parquet_parallel
 
 run_parallel = steps_to_parquet_parallel.run_parallel
 resolve_destination = steps_to_parquet_parallel.resolve_destination

tests/test_train.py

@@ -5,6 +5,7 @@ import csv
 import math
 import tempfile
 from pathlib import Path
+from unittest.mock import MagicMock
 
 import pytest
 import torch
@@ -17,6 +18,7 @@ from giant.constants import (
     SEC_SLOT_DIM,
     X_DIM,
 )
+from giant.data.dataset import StepBatch
 from giant.model.network import build_critics, build_models
 from giant.training import (
     FlowDDPMStageTrainer,
@@ -316,7 +318,7 @@ def _fake_batches(n_batches, batch_size, seed=0):
         sec_cont = torch.randn(batch_size, K_MAX, SEC_SLOT_DIM, generator=g)
         proc_idx = torch.zeros(batch_size, dtype=torch.long)
         sec_type_idx = torch.zeros(batch_size, K_MAX, dtype=torch.long)
-        batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx))
+        batches.append(StepBatch(cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx))
     return batches
 
 
@@ -586,6 +588,60 @@ def test_stage_spec_from_config_omitted_decoder_and_particle_type_match_default_
     assert spec.particle_type.target == "onehot"
 
 
+def _routed_stage1_trainer(lambda_balance, lambda_proc, lambda_entropy):
+    cfg = _base_cfg()
+    cfg["stage1_model"]["router"] = {
+        "enabled": True,
+        "type": "energy",
+        "n_experts": 3,
+        "temperature": 0.5,
+        "learn_centers": True,
+        "lambda_balance": lambda_balance,
+        "lambda_proc": lambda_proc,
+        "lambda_entropy": lambda_entropy,
+    }
+    model_config = _model_config(cfg)
+    models = build_models(model_config)
+    critics = build_critics(model_config)
+    trainers = build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4)
+    return trainers["stage1"]
+
+
+def test_router_aux_losses_skipped_when_lambda_zero_but_run_when_positive():
+    """Gitea #31: FlowDDPMStageTrainer._compute must not call
+    router.balance_loss/classify_loss/entropy_loss when the corresponding
+    lambda is 0 (the default) -- those calls do their own router.gate(...)
+    forward pass that is wasted once the term is masked out of the total
+    loss anyway. Checked both ways: zero lambdas must skip all three calls,
+    positive lambdas must still make them (the guard must not accidentally
+    suppress the real path)."""
+    batch = _fake_batches(1, 4)[0]
+    device = torch.device("cpu")
+
+    trainer_zero = _routed_stage1_trainer(0.0, 0.0, 0.0)
+    router_zero = trainer_zero.router
+    router_zero.balance_loss = MagicMock(wraps=router_zero.balance_loss)
+    router_zero.classify_loss = MagicMock(wraps=router_zero.classify_loss)
+    router_zero.entropy_loss = MagicMock(wraps=router_zero.entropy_loss)
+    stats_zero = trainer_zero.step(batch, device, global_step=1)
+    assert router_zero.balance_loss.call_count == 0
+    assert router_zero.classify_loss.call_count == 0
+    assert router_zero.entropy_loss.call_count == 0
+    assert stats_zero["loss_balance"] == 0.0
+    assert stats_zero["loss_proc"] == 0.0
+    assert stats_zero["loss_entropy"] == 0.0
+
+    trainer_pos = _routed_stage1_trainer(0.1, 0.1, 0.01)
+    router_pos = trainer_pos.router
+    router_pos.balance_loss = MagicMock(wraps=router_pos.balance_loss)
+    router_pos.classify_loss = MagicMock(wraps=router_pos.classify_loss)
+    router_pos.entropy_loss = MagicMock(wraps=router_pos.entropy_loss)
+    trainer_pos.step(batch, device, global_step=1)
+    assert router_pos.balance_loss.call_count == 1
+    assert router_pos.classify_loss.call_count == 1
+    assert router_pos.entropy_loss.call_count == 1
+
+
 # --- AR trainer wiring (v0.3.0 step 5) --------------------------------------
 
 

tests/test_validate.py

@@ -2,6 +2,7 @@ import numpy as np
 import torch
 
 from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM
+from giant.data.dataset import StepBatch
 from giant.model.network import Stage1Model, Stage2OneShot, stage2_trunk_sec_dim
 from giant.validate import validate_marginals
 
@@ -46,8 +47,7 @@ def _tiny_models(particle_type_cfg: dict | None = None):
 
 
 def _loader(B: int = 4, n_batches: int = 2, n_sec_value: int = 0, n_classes: int = 8):
-    """A val_loader matching StreamingStepsDataset's 7-tuple batch shape:
-    (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, sec_type_idx)."""
+    """A val_loader matching StreamingStepsDataset's StepBatch shape."""
     batches = []
     for _ in range(n_batches):
         cond_cont = torch.randn(B, COND_DIM)
@@ -57,7 +57,7 @@ def _loader(B: int = 4, n_batches: int = 2, n_sec_value: int = 0, n_classes: int
         sec_cont = torch.randn(B, _K_MAX, SEC_SLOT_DIM)
         proc_idx = torch.zeros(B, dtype=torch.long)
         sec_type_idx = torch.randint(0, n_classes, (B, _K_MAX), dtype=torch.long)
-        batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx))
+        batches.append(StepBatch(cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx))
     return batches