Merge pull request 'Fix/issue 40' (#64) from fix/issue-40 into master

Reviewed-on: #64

configs/baseline.toml

@@ -0,0 +1,129 @@
+# GIANT reference baseline (v0.3 schema).
+#
+# The fixed comparison point every future architecture variant is measured
+# against. Chosen so that each experimental axis the roadmap cares about
+# (routed trunk, WGAN generators, attention history, shared conditioning,
+# embedding/onehot conditioning) is a *single* edit away from this file.
+#
+# Rationale for the choices below, from the runs already on record
+# (analysis_runs/ + the `giant` W&B project):
+#
+#   * flow, not wgan, for both stages. Ranking the five existing rollouts by
+#     mean Jensen-Shannon divergence against the Geant4 reference, the plain
+#     non-routed flow model wins (0.172) over the routed flow runs
+#     (0.197/0.200) and both WGAN runs (0.218/0.234) — and it beats them by
+#     ~7x on per-event total deposited energy and by 3-10x on every
+#     per-PDG marginal. WGAN stays a variant, not the reference.
+#
+#   * no router. The routed runs are not better, and soft-mixing 10 small
+#     experts costs ~10x per-pass throughput at train time (29k samples/s vs
+#     the WGAN runs' 52-116k), which is what made those runs take ~110 h for
+#     30 epochs.
+#
+#   * hidden_dim 512 / 6 blocks per stage. The best-scoring rollout so far
+#     was hidden_dim 1024, but at 4x the trunk FLOPs of 512. 512/6 sits in
+#     the same weight class as the variants it will be compared against and
+#     leaves headroom to train it properly rather than cheaply.
+#
+#   * dropout 0.0. Training set is ~5e8 steps against <1e7 parameters;
+#     capacity overfitting is not the binding constraint, and every recent
+#     run used 0.0.
+#
+# Known weak spots this baseline is expected to *exhibit* (they are the
+# reason for the comparisons, not a reason to retune this file): every model
+# on record under-produces steps per event by ~2x (rollout ~7e4 vs Geant4
+# ~1.4e5) and secondaries per event by 2-3.5x (~2-3e4 vs 7.2e4), and n_sec
+# head accuracy sits at 0.863-0.867 regardless of size or objective.
+
+[meta]
+# REQUIRED. Without it config.migrate_config reads this file as v0.2 and
+# rewrites it from V02_FIXED_FACTS — silently forcing decoder = "one_shot",
+# particle_type.target = "physical" and the v0.2 default sizes, while still
+# passing validate_config.
+config_version = 3
+
+[conditioning]
+# Physical-property MLPs rather than learned vocab embeddings: computable for
+# any PDG code / material, which is what the held-out-species and
+# held-out-material generalization comparisons need.
+out_dim      = 128
+share_stages = false
+
+# n_layers = 2 rather than the v0.3 default of 1: v0.2's conditioning MLP was
+# always 2 deep (see _migration.V02_FIXED_FACTS), so this keeps the encoder
+# identical to the architecture that produced the results cited above.
+[conditioning.particle]
+type     = "physical"
+emb_dim  = 16
+n_layers = 2
+
+[conditioning.material]
+type     = "physical"
+emb_dim  = 16
+n_layers = 2
+
+[stage1_model]
+generator    = "flow"
+hidden_dim   = 512
+n_res_blocks = 6
+dropout      = 0.0
+
+[stage2_model]
+# The v0.3 pivot: autoregressive in descending-energy order with a
+# categorical species target, which is the agreed response to the 2026-08-03
+# secondary-species failure. Flow (not the schema default wgan) so the
+# baseline varies only the decoder relative to the best v0.2 result.
+#
+# COST, measured (RTX 4070, bs 4096, 10 ODE steps), not estimated:
+# sample.sample_secondaries_ar loops `for k in range(k_max)` unconditionally
+# — all 15 slots regardless of predicted n_sec — so a flow AR token costs
+# k_max * steps = 150 stage-2 calls per physics step. That makes this block
+# the dominant cost on both sides:
+#   training   flow AR 29.5k samp/s  vs flow one-shot 190.7k samp/s  (6.5x)
+#   inference  flow AR  8.5k step/s  vs flow one-shot  68.7k step/s  (8.1x)
+# Accepted deliberately: one-shot is the configuration whose secondary
+# species distribution failed, and that failure is what v0.3 exists to fix.
+decoder      = "autoregressive"
+generator    = "flow"
+hidden_dim   = 512
+n_res_blocks = 6
+dropout      = 0.0
+k_max        = 15
+
+[stage2_model.autoregressive]
+history         = "markov"
+teacher_forcing = "always"
+
+[stage2_model.particle_type]
+target = "onehot"
+# Decoupled from conditioning.particle.emb_dim (gitea #29). 32 classes + the
+# "other" bucket keeps essentially all real secondary species out of "other"
+# without making the head expensive.
+n_classes    = 32
+other_policy = "sample"
+
+[train]
+epochs        = 50
+# Sized for ONE NVIDIA L40S on deepthought2 (46068 MiB; the box has two, and
+# CLAUDE.md's shared-machine rule allows a single GPU). From a measured
+# linear fit of this exact config's training step on the local RTX 4070:
+#   peak reserved MiB = 0.9736 * batch_size + 115
+# so 36864 reserves ~36.0 GiB, i.e. 78% of the card, leaving ~10 GiB of
+# headroom for fragmentation and the CUDA context. Throughput is already
+# flat above bs~4096 on the 4070, so this is chosen for occupancy on the
+# larger card, not for step efficiency — and it sits next to the 43008/32768
+# of the runs lr = 3e-4 was proven at.
+batch_size    = 36864
+lr            = 3e-4
+warmup_epochs = 3
+weight_decay  = 0.01
+ema_decay     = 0.9999
+val_fraction  = 0.1
+num_workers   = 4
+seed          = 0
+# The marginal/KL pass is expensive (~5000 s on top of an epoch), so keep it
+# to every 10th epoch; the cheap per-epoch val loss still runs every epoch.
+validate_every = 10
+validate_steps = 10
+wandb          = true
+wandb_project  = "giant"

giant/analysis/reduce.py

@@ -29,8 +29,17 @@ from giant.constants import TERM_ESCAPED
 
 
 def _bin_expr(value: pl.Expr, lo: float, hi: float, nbins: int) -> pl.Expr:
-    """Uniform bin index of ``value`` over ``[lo, hi]`` into ``nbins`` bins."""
-    return ((value - lo) / (hi - lo) * nbins).floor().cast(pl.Int32).clip(0, nbins - 1)
+    """Uniform bin index of ``value`` over ``[lo, hi]`` into ``nbins`` bins.
+
+    Out-of-range values clamp into the edge bins, and the clamp deliberately
+    happens in f64 *before* the integer cast: a rollout is free to emit a wildly
+    out-of-range outlier (a step_length of 1e10 mm, say) or an inf, whose
+    unclamped bin index overflows i32 and makes the cast fail outright. NaN has
+    no edge to clamp to, so it becomes null and is dropped by the callers below
+    — the same thing ``np.histogram`` does with it.
+    """
+    idx = ((value - lo) / (hi - lo) * nbins).floor().clip(0, nbins - 1)
+    return pl.when(idx.is_nan()).then(None).otherwise(idx).cast(pl.Int32)
 
 
 def hist1d(
@@ -50,6 +59,7 @@ def hist1d(
     group = pl.lit(0, dtype=pl.Int64) if group is None else group
     res = (
         lf.select(group.alias("_g"), _bin_expr(value, lo, hi, nbins).alias("_b"))
+        .drop_nulls("_b")
         .group_by("_g", "_b")
         .agg(pl.len().alias("_n"))
         .collect(engine="streaming")
@@ -188,6 +198,7 @@ def profile_partial(
             _bin_expr(coord, lo, hi, nbins).alias("_b"),
             weight.alias("_w"),
         )
+        .drop_nulls("_b")
         .group_by("event_id", "_b")
         .agg(pl.col("_w").sum().alias("_ws"))
         .collect(engine="streaming")

giant/cli.py

@@ -1043,8 +1043,12 @@ def predict(
             # A fresh v0.3.0 Stage1Model owns no n_sec_head —
             # sample_stage1 returns n_sec_pred=None then, so ask stage 2.
             n_sec_pred = resolve_n_sec(model, sec_decoder, cc, ck, stage1_norm, n_sec_pred)
-            sec_cont, sec_type, _sec_valid_pred = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
-            n_sec_pred_np = n_sec_pred.cpu().numpy()
+            sec_cont, sec_type, sec_valid_pred = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
+            # A stop-token decoder resolves n_sec_pred=None above — read the
+            # real count back off sec_valid_pred instead (a no-op round trip
+            # under every other n_sec.mode, where sec_valid_pred was built
+            # FROM n_sec_pred in the first place).
+            n_sec_pred_np = sec_valid_pred.sum(dim=-1).cpu().numpy()
 
         pred = stage1_norm.cpu().numpy()  # normalised
 

giant/cond_layout.py

@@ -0,0 +1,103 @@
+"""Single source of truth for the conditioning arrays' column layout (gitea #37).
+
+`cond_cont` and `cond_cat` are built in `giant.data.transforms` and consumed in
+`giant.model.encoders` / `giant.model.routers`. Their column order used to be
+written down independently on each side, kept in sync only by parallel comments
+— so getting it wrong produced silently mis-indexed columns rather than an
+exception, and adding a conditioning axis meant a coordinated multi-file edit.
+
+`CondLayout` owns that order. Both sides construct one from the same
+`conditioning.particle.type` / `conditioning.material.type` pair and read named
+slices off it, so the layout is stated exactly once. This module depends only on
+`giant.constants`, so both the data and model packages can import it.
+"""
+
+from dataclasses import dataclass
+from typing import ClassVar
+
+from giant.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM
+
+# The three per-axis conditioning modes. Mirrors giant.config.Conditioning,
+# which this module deliberately does not import (giant.config pulls in the
+# whole model package).
+AXIS_TYPES = ("physical", "embedding", "onehot")
+
+
+@dataclass(frozen=True)
+class CondLayout:
+    """Column layout of `cond_cont`/`cond_cat` for one (particle, material) mode pair.
+
+    `cond_cont` is unconditionally `COND_DIM` wide regardless of mode: the base
+    block, then the particle physical block, then the material physical block.
+    An axis that isn't `"physical"` gets its block zero-filled and never reads
+    it (see `giant.data.transforms._physical_cond_columns`), so the widths are
+    mode-independent and only the *meaning* of a block changes.
+
+    `cond_cat` is 2 to 4 wide. Columns `PDG_COL`/`MAT_COL` are always the dense
+    training-vocab index; an axis in `"onehot"` mode appends one more column
+    holding its top-N-plus-other class index, particle before material.
+    """
+
+    particle_type: str
+    material_type: str
+
+    # cond_cat's dense-vocab columns, present in every mode. Under
+    # "physical"/"onehot" they are a reporting/router convenience the
+    # ConditionEncoder never reads; under "embedding" they are the signal.
+    PDG_COL: ClassVar[int] = 0
+    MAT_COL: ClassVar[int] = 1
+
+    def __post_init__(self) -> None:
+        if self.particle_type not in AXIS_TYPES:
+            raise ValueError(f"unknown conditioning.particle.type {self.particle_type!r}")
+        if self.material_type not in AXIS_TYPES:
+            raise ValueError(f"unknown conditioning.material.type {self.material_type!r}")
+
+    @classmethod
+    def from_types(cls, particle_type: str, material_type: str) -> "CondLayout":
+        """Named constructor — the entry point both sides use."""
+        return cls(particle_type=particle_type, material_type=material_type)
+
+    # --- cond_cont ---------------------------------------------------------
+
+    @property
+    def base(self) -> slice:
+        """pre_pos(3), log(pre_E)(1), pre_dir(3), layer_id(1)."""
+        return slice(0, COND_DIM_BASE)
+
+    @property
+    def particle_phys(self) -> slice:
+        """log(mass), charge — see `giant.particles`."""
+        return slice(COND_DIM_BASE, COND_DIM_BASE + PARTICLE_PHYS_DIM)
+
+    @property
+    def material_phys(self) -> slice:
+        """Z_eff, A_eff, log(density), log(X0), log(lambda_int) — see `giant.materials`."""
+        start = COND_DIM_BASE + PARTICLE_PHYS_DIM
+        return slice(start, start + MATERIAL_PHYS_DIM)
+
+    @property
+    def cont_dim(self) -> int:
+        return COND_DIM
+
+    # --- cond_cat ----------------------------------------------------------
+
+    @property
+    def particle_topn_col(self) -> int | None:
+        """Column of the particle top-N class index, or `None` if not `"onehot"`."""
+        return self.MAT_COL + 1 if self.particle_type == "onehot" else None
+
+    @property
+    def material_topn_col(self) -> int | None:
+        """Column of the material top-N class index, or `None` if not `"onehot"`.
+
+        Comes after the particle top-N column when both axes are `"onehot"`.
+        """
+        if self.material_type != "onehot":
+            return None
+        return self.MAT_COL + (2 if self.particle_type == "onehot" else 1)
+
+    @property
+    def cat_dim(self) -> int:
+        """Total `cond_cat` width: 2, plus one column per `"onehot"` axis."""
+        return self.MAT_COL + 1 + (self.particle_type == "onehot") + (self.material_type == "onehot")

giant/config.py

@@ -15,6 +15,7 @@ import numpy as np
 import torch
 
 from giant._migration import V02_FIXED_FACTS, V02_MODEL_KEY_TO_STAGES, reject_legacy_router_expert_sizing
+from giant.model.history import HISTORY_REGISTRY
 
 
 class Conditioning(str, Enum):
@@ -336,6 +337,35 @@ class RouterConfig:
         }
 
 
+@dataclass(frozen=True)
+class TrunkConfig:
+    """`stage1_model.trunk`/`stage2_model.trunk`: selects the trunk's expert
+    *body* architecture from `giant.model.trunks.TRUNK_REGISTRY` (default
+    `"resmlp"` — today's only body, `input_proj -> ResBlock stack ->
+    out_proj`). Orthogonal to whether that body is mixed: mixing is still
+    controlled entirely by `router.enabled`/`router.n_experts` on the same
+    stage, unaffected by this block. A future body's own hyperparameters
+    (e.g. a transformer's `n_heads`/`n_layers`) would get their own sibling
+    field here, matching how `flow`/`ddpm`/`wgan` already coexist selected by
+    `generator`.
+
+    `block_conditioning` selects each body's conditioning-injection mechanism
+    from `giant.model.layers.BLOCK_REGISTRY` — `"add"` (default, today's
+    conditional-bias `ResBlock`, bit-identical to pre-gitea-#34 behaviour),
+    `"film"`, or `"adaln"`."""
+
+    type: str = "resmlp"
+    block_conditioning: str = "add"
+
+    @classmethod
+    def from_dict(cls, d: dict | None) -> "TrunkConfig":
+        d = d or {}
+        return cls(type=d.get("type", "resmlp"), block_conditioning=d.get("block_conditioning", "add"))
+
+    def to_dict(self) -> dict:
+        return {"type": self.type, "block_conditioning": self.block_conditioning}
+
+
 @dataclass(frozen=True)
 class Stage2RouterConfig(RouterConfig):
     # true: stage 2 shares stage 1's Router module instance, so expert i in
@@ -378,18 +408,27 @@ class Stage2RouterConfig(RouterConfig):
 class NSecConfig:
     # "head": a classifier over {0..k_max} on the condition encoding alone
     # (no diffusion noise), callable independently at inference.
-    # "stop_token": an EOS-style implicit stop — accepted by the schema but
-    # not implemented in v0.3.0 (see validate_config).
+    # "stop_token": an EOS-style per-slot stop head on the autoregressive
+    # secondary decoder (Stage2Autoregressive only — see validate_config),
+    # evaluated against the generated prefix instead of conditioning alone.
+    # Replaces n_sec_head entirely: the two are mutually exclusive, so this
+    # mode builds no n_sec_head and stage2_model.n_sec.lambda instead weights
+    # the stop head's BCE term.
     # "truth": take n_sec from ground truth — standalone stage-2 evaluation
     # only, never for rollout.
     mode: str = "head"
-    lambda_weight: float = 0.1  # dict key "lambda" — cross-entropy weight for the head
+    lambda_weight: float = 0.1  # dict key "lambda" — cross-entropy/BCE weight for the head
     # Which stage's module physically owns the n_sec_head weights: "stage2" (default,
     # fresh v0.3.0 runs — Stage2OneShot/Stage2Autoregressive builds it) or "stage1"
     # (a migrated v0.2 checkpoint — see network._migrate_legacy_model_config, whose
     # n_sec head was trained against Stage 1's own ConditionEncoder output and so has
     # to stay attached there, not just be labeled as such).
     owner: str = "stage2"
+    # mode="stop_token" only: how sample_secondaries_ar turns a slot's stop logit into a
+    # stop/continue decision. "greedy": sigmoid(logit) >= 0.5 (deterministic). "sample":
+    # a Bernoulli draw at sigmoid(logit) (a real sample from the learned length
+    # distribution, at the cost of an extra RNG draw per slot).
+    stop_sampling: str = "greedy"
 
     @classmethod
     def from_dict(cls, d: dict | None) -> "NSecConfig":
@@ -398,10 +437,16 @@ class NSecConfig:
             mode=d.get("mode", "head"),
             lambda_weight=d.get("lambda", 0.1),
             owner=d.get("owner", "stage2"),
+            stop_sampling=d.get("stop_sampling", "greedy"),
         )
 
     def to_dict(self) -> dict:
-        return {"mode": self.mode, "lambda": self.lambda_weight, "owner": self.owner}
+        return {
+            "mode": self.mode,
+            "lambda": self.lambda_weight,
+            "owner": self.owner,
+            "stop_sampling": self.stop_sampling,
+        }
 
 
 @dataclass(frozen=True)
@@ -486,6 +531,64 @@ class AutoregressiveConfig:
         }
 
 
+@dataclass(frozen=True)
+class HeadConfig:
+    """A single classifier head's shape — `n_sec_head`/`type_head` (gitea
+    #36 deduplicated their five identical hand-rolled
+    `Linear -> SiLU -> Linear` definitions into
+    `giant.model.layers.build_mlp_head`, which this config drives).
+    `hidden_ratio=0.5`/`depth=2` are the exact pre-#36 hardcoded values
+    (hidden width = `hidden_dim // 2`, one hidden layer), so omitting a
+    `heads` block — including every migrated v0.2 config — reproduces the
+    old architecture bit-for-bit."""
+
+    hidden_ratio: float = 0.5  # hidden width = round(hidden_dim * hidden_ratio)
+    depth: int = 2  # matches build_mlp_head's depth
+
+    @classmethod
+    def from_dict(cls, d: dict | None) -> "HeadConfig":
+        d = d or {}
+        return cls(hidden_ratio=d.get("hidden_ratio", 0.5), depth=d.get("depth", 2))
+
+    def to_dict(self) -> dict:
+        return {"hidden_ratio": self.hidden_ratio, "depth": self.depth}
+
+
+@dataclass(frozen=True)
+class Stage1HeadsConfig:
+    """Stage 1 only ever owns `n_sec_head`, and only for a migrated v0.2
+    checkpoint (`stage2_model.n_sec.owner = "stage1"`) — see
+    `Stage1Model`'s docstring."""
+
+    n_sec: HeadConfig = field(default_factory=HeadConfig)
+
+    @classmethod
+    def from_dict(cls, d: dict | None) -> "Stage1HeadsConfig":
+        d = d or {}
+        return cls(n_sec=HeadConfig.from_dict(d.get("n_sec")))
+
+    def to_dict(self) -> dict:
+        return {"n_sec": self.n_sec.to_dict()}
+
+
+@dataclass(frozen=True)
+class Stage2HeadsConfig:
+    """`n_sec` and `type` are independently configurable — n_sec accuracy
+    and secondary-species accuracy are separately known weak spots (gitea
+    #36)."""
+
+    n_sec: HeadConfig = field(default_factory=HeadConfig)
+    type: HeadConfig = field(default_factory=HeadConfig)
+
+    @classmethod
+    def from_dict(cls, d: dict | None) -> "Stage2HeadsConfig":
+        d = d or {}
+        return cls(n_sec=HeadConfig.from_dict(d.get("n_sec")), type=HeadConfig.from_dict(d.get("type")))
+
+    def to_dict(self) -> dict:
+        return {"n_sec": self.n_sec.to_dict(), "type": self.type.to_dict()}
+
+
 @dataclass(frozen=True)
 class Stage1ModelConfig:
     # false skips building/training stage 1 entirely. The resulting
@@ -509,6 +612,8 @@ class Stage1ModelConfig:
     ddpm: DdpmConfig = field(default_factory=DdpmConfig)
     wgan: Stage1WganConfig = field(default_factory=Stage1WganConfig)
     router: RouterConfig = field(default_factory=RouterConfig)
+    trunk: TrunkConfig = field(default_factory=TrunkConfig)
+    heads: Stage1HeadsConfig = field(default_factory=Stage1HeadsConfig)
 
     @classmethod
     def from_dict(cls, d: dict | None) -> "Stage1ModelConfig":
@@ -524,6 +629,8 @@ class Stage1ModelConfig:
             ddpm=DdpmConfig.from_dict(d.get("ddpm")),
             wgan=Stage1WganConfig.from_dict(d.get("wgan")),
             router=RouterConfig.from_dict(d.get("router")),
+            trunk=TrunkConfig.from_dict(d.get("trunk")),
+            heads=Stage1HeadsConfig.from_dict(d.get("heads")),
         )
 
     def to_dict(self) -> dict:
@@ -538,6 +645,8 @@ class Stage1ModelConfig:
             "ddpm": self.ddpm.to_dict(),
             "wgan": self.wgan.to_dict(),
             "router": self.router.to_dict(),
+            "trunk": self.trunk.to_dict(),
+            "heads": self.heads.to_dict(),
         }
 
 
@@ -575,6 +684,8 @@ class Stage2ModelConfig:
     ddpm: DdpmConfig = field(default_factory=DdpmConfig)
     wgan: Stage2WganConfig = field(default_factory=Stage2WganConfig)
     router: Stage2RouterConfig = field(default_factory=Stage2RouterConfig)
+    trunk: TrunkConfig = field(default_factory=TrunkConfig)
+    heads: Stage2HeadsConfig = field(default_factory=Stage2HeadsConfig)
 
     @classmethod
     def from_dict(cls, d: dict | None) -> "Stage2ModelConfig":
@@ -597,6 +708,8 @@ class Stage2ModelConfig:
             ddpm=DdpmConfig.from_dict(d.get("ddpm")),
             wgan=Stage2WganConfig.from_dict(d.get("wgan")),
             router=Stage2RouterConfig.from_dict(d.get("router")),
+            trunk=TrunkConfig.from_dict(d.get("trunk")),
+            heads=Stage2HeadsConfig.from_dict(d.get("heads")),
         )
 
     def to_dict(self) -> dict:
@@ -618,6 +731,8 @@ class Stage2ModelConfig:
             "ddpm": self.ddpm.to_dict(),
             "wgan": self.wgan.to_dict(),
             "router": self.router.to_dict(),
+            "trunk": self.trunk.to_dict(),
+            "heads": self.heads.to_dict(),
         }
 
 
@@ -1268,11 +1383,23 @@ def validate_config(cfg: dict) -> None:
         )
 
     if _get_path(cfg, "stage2_model.n_sec.mode") == "stop_token":
-        raise ValueError(
-            "stage2_model.n_sec.mode = 'stop_token' is accepted by the schema "
-            "but not implemented in v0.3.0 — use 'head' (default) or 'truth' "
-            "(standalone stage-2 evaluation only, never for rollout)"
-        )
+        if _get_path(cfg, "stage2_model.decoder") != "autoregressive":
+            raise ValueError(
+                "stage2_model.n_sec.mode = 'stop_token' requires "
+                "stage2_model.decoder = 'autoregressive' — there is no "
+                "per-token loop to stop under 'one_shot'"
+            )
+        if _get_path(cfg, "stage2_model.n_sec.owner") != "stage2":
+            raise ValueError(
+                "stage2_model.n_sec.mode = 'stop_token' requires "
+                "stage2_model.n_sec.owner = 'stage2' — a migrated v0.2 "
+                "checkpoint's stage-1 n_sec_head has no per-token "
+                "conditioning to hang an EOS decision off"
+            )
+
+    stop_sampling = _get_path(cfg, "stage2_model.n_sec.stop_sampling")
+    if stop_sampling not in ("greedy", "sample"):
+        raise ValueError(f"stage2_model.n_sec.stop_sampling = {stop_sampling!r} — must be 'greedy' or 'sample'")
 
     if _get_path(cfg, "stage2_model.stage1_context") == "sampled":
         raise ValueError(
@@ -1306,8 +1433,10 @@ def validate_config(cfg: dict) -> None:
                 "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'")
+        if history not in HISTORY_REGISTRY:
+            raise ValueError(
+                f"stage2_model.autoregressive.history = {history!r} — must be one of {sorted(HISTORY_REGISTRY)}"
+            )
         teacher_forcing = _get_path(cfg, "stage2_model.autoregressive.teacher_forcing")
         if teacher_forcing not in ("always", "scheduled", "never"):
             raise ValueError(
@@ -1400,6 +1529,10 @@ _OUT_DIR_NAME_CANDIDATES = [
         "particle_type_target",
         _path_candidate("stage2_model.particle_type.target", "pt-"),
     ),
+    ("stage1_trunk_type", _path_candidate("stage1_model.trunk.type", "s1t-")),
+    ("stage2_trunk_type", _path_candidate("stage2_model.trunk.type", "s2t-")),
+    ("stage1_block_cond", _path_candidate("stage1_model.trunk.block_conditioning", "s1bc-")),
+    ("stage2_block_cond", _path_candidate("stage2_model.trunk.block_conditioning", "s2bc-")),
     ("stage1_router", _router_candidate("stage1_model", "s1")),
     ("stage2_router", _router_candidate("stage2_model", "s2")),
     (

giant/data/transforms.py

@@ -3,6 +3,7 @@ from typing import NamedTuple
 
 import numpy as np
 
+from giant.cond_layout import CondLayout
 from giant.constants import K_MAX
 
 _EPS = 1e-8
@@ -13,6 +14,14 @@ _EPS = 1e-8
 # the conservation it slightly softens is physically negligible (~0.001%).
 _SIMPLEX_FLOOR = 1e-5
 
+# Upper clip for a raw predicted log_mass before inv_log_transform: exp(y)
+# must stay well inside float32 range (~3.4e38, i.e. y < ~88.7) or it
+# overflows to inf, which — like the negative-mass case below — blows up the
+# next log_transform call once that mass is fed back in as conditioning.
+# 80.0 leaves comfortable headroom while still being far beyond any physical
+# particle mass a converged model would ever predict.
+_LOG_MASS_MAX = 80.0
+
 
 def log_transform(x: np.ndarray, eps: float = _EPS) -> np.ndarray:
     x = np.asarray(x, dtype=np.float32)
@@ -684,27 +693,31 @@ def decode_secondaries(
     log_mass = sec_cont[:, :, 4]  # (N, K)
     charge = sec_cont[:, :, 5]  # (N, K)
 
-    # mass is non-negative by construction (inv_log_transform of a real
-    # number is always > 0); clip to 0 for padded/invalid slots rather than
-    # leaving a spurious small positive floor from the log inverse.
+    # log_mass is a raw model prediction, not itself the output of
+    # log_transform, so it can land far outside the range that round-trips
+    # cleanly through inv_log_transform: too negative and exp(log_mass)
+    # undershoots _EPS, making inv_log_transform go slightly negative; too
+    # positive and exp(log_mass) overflows float32 to inf. Either one then
+    # blows up the next log_transform call on this track's mass once it's
+    # fed back in as conditioning for a further rollout step
+    # (giant/rollout.py -> build_cond_features -> _physical_cond_columns).
+    # Clip to a range whose inverse is guaranteed finite and >= 0 before
+    # that can happen; clip to 0 separately for padded/invalid slots rather
+    # than leaving a spurious small positive floor.
+    log_mass = np.clip(log_mass, np.log(_EPS), _LOG_MASS_MAX)
     sec_mass = np.where(sec_valid, inv_log_transform(log_mass), 0.0).astype(np.float32)
     sec_charge = np.where(sec_valid, charge, 0.0).astype(np.float32)
 
     return sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid
 
 
-def _physical_cond_columns(
-    data: dict[str, np.ndarray],
-    particle_conditioning: str,
-    material_conditioning: str,
-) -> np.ndarray:
+def _physical_cond_columns(data: dict[str, np.ndarray], layout: CondLayout) -> np.ndarray:
     """(N, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM) physical conditioning columns.
 
     The particle and material blocks are gated independently and may mix
     freely — e.g. material `physical` with particle `embedding` — so e.g.
-    `particle_conditioning="embedding"` + `material_conditioning="physical"`
-    zero-fills only the particle columns and computes the material ones for
-    real.
+    `particle_type="embedding"` + `material_type="physical"` zero-fills only
+    the particle columns and computes the material ones for real.
 
     "embedding"/"onehot" zero-fill their block (cheap, and ConditionEncoder
     never reads these columns in either mode — so an unfilled
@@ -721,7 +734,7 @@ def _physical_cond_columns(
 
     n = len(next(iter(data.values())))
 
-    if particle_conditioning == "physical":
+    if layout.particle_type == "physical":
         from giant.particles import particle_phys_array
 
         if "mass" in data and "charge" in data:
@@ -730,12 +743,10 @@ def _physical_cond_columns(
         else:
             mass, charge = particle_phys_array(data["pdg"]).T
         particle_cols = np.column_stack([log_transform(mass), charge])
-    elif particle_conditioning in ("embedding", "onehot"):
-        particle_cols = np.zeros((n, PARTICLE_PHYS_DIM), dtype=np.float32)
     else:
-        raise ValueError(f"unknown conditioning.particle.type {particle_conditioning!r}")
+        particle_cols = np.zeros((n, PARTICLE_PHYS_DIM), dtype=np.float32)
 
-    if material_conditioning == "physical":
+    if layout.material_type == "physical":
         from giant.materials import material_properties_array
 
         z_eff, a_eff, density, x0, lambda_int = material_properties_array(data["material"]).T
@@ -748,14 +759,66 @@ def _physical_cond_columns(
                 log_transform(lambda_int),
             ]
         )
-    elif material_conditioning in ("embedding", "onehot"):
-        material_cols = np.zeros((n, MATERIAL_PHYS_DIM), dtype=np.float32)
     else:
-        raise ValueError(f"unknown conditioning.material.type {material_conditioning!r}")
+        material_cols = np.zeros((n, MATERIAL_PHYS_DIM), dtype=np.float32)
 
     return np.column_stack([particle_cols, material_cols]).astype(np.float32)
 
 
+def _build_cond_arrays(
+    data: dict[str, np.ndarray],
+    pdg_map: dict[int, int],
+    mat_map: dict[str, int],
+    layout: CondLayout,
+    pdg_topn_map: dict[int, int] | None,
+    mat_topn_map: dict[str, int] | None,
+) -> tuple[np.ndarray, np.ndarray]:
+    """The un-normalized `(cond_cont, cond_cat)` pair, in `layout`'s column order.
+
+    Both `build_cond_features` and `build_features` go through here, so the
+    column order — and everything that depends on it — is stated once. See
+    `giant.cond_layout.CondLayout` for the layout itself.
+    """
+    cond_cont = np.column_stack(
+        [
+            data["pre_pos"],
+            log_transform(data["pre_E"]),
+            data["pre_dir"],
+            data["layer_id"].astype(np.float32),
+        ]
+    ).astype(np.float32)  # (N, COND_DIM_BASE=8)
+    cond_cont = np.column_stack([cond_cont, _physical_cond_columns(data, layout)]).astype(
+        np.float32
+    )  # (N, COND_DIM=15)
+
+    # In "physical" mode cond_cat's first two columns are only a
+    # reporting/router convenience — ConditionEncoder never reads them
+    # (giant/model/encoders.py) — so a species/material outside the training
+    # vocab (the whole point of physical-property conditioning) gets a dummy
+    # index instead of raising. In "embedding" mode those columns ARE the
+    # conditioning signal, so an unmapped value must still raise loudly
+    # rather than silently misassign. In "onehot" mode they again go unread
+    # (the topN columns below are the real signal), so they're as permissive
+    # as "physical". Each axis's strictness is independent.
+    pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=layout.particle_type == "embedding")
+    mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=layout.material_type == "embedding")
+    # Which extra columns exist is the layout's call, not "did the caller
+    # happen to pass a map" — that's what used to let the producer and
+    # ConditionEncoder disagree. A map for a non-"onehot" axis is unused.
+    cat_cols = [pdg_idx, mat_idx]
+    if layout.particle_topn_col is not None:
+        if pdg_topn_map is None:
+            raise ValueError("conditioning.particle.type='onehot' needs pdg_topn_map")
+        cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map))
+    if layout.material_topn_col is not None:
+        if mat_topn_map is None:
+            raise ValueError("conditioning.material.type='onehot' needs mat_topn_map")
+        cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map))
+    cond_cat = np.column_stack(cat_cols)  # (N, layout.cat_dim)
+
+    return cond_cont, cond_cat
+
+
 def build_cond_features(
     data: dict[str, np.ndarray],
     pdg_map: dict[int, int],
@@ -773,49 +836,16 @@ def build_cond_features(
     `material_conditioning="physical"` is a valid mix.
 
     `pdg_topn_map`/`mat_topn_map` (a top-N-plus-other `class_map`, see
-    `giant.data.loader.build_topn_map_from_files`) append extra `cond_cat`
-    columns read by `ConditionEncoder`'s `"onehot"` mode: pdg topN index at
-    column 2 (iff `pdg_topn_map` given), material topN index at column 3
-    (iff `mat_topn_map` given, after column 2 if both are). Only ever given when
-    the corresponding axis is `"onehot"`; `cond_cat` stays `(N, 2)` otherwise.
+    `giant.data.loader.build_topn_map_from_files`) supply the extra `cond_cat`
+    columns read by `ConditionEncoder`'s `"onehot"` mode, and are required
+    whenever the corresponding axis is `"onehot"`. See
+    `giant.cond_layout.CondLayout` for which columns exist where.
     """
-    cond_cont = np.column_stack(
-        [
-            data["pre_pos"],
-            log_transform(data["pre_E"]),
-            data["pre_dir"],
-            data["layer_id"].astype(np.float32),
-        ]
-    ).astype(np.float32)
-    cond_cont = np.column_stack(
-        [
-            cond_cont,
-            _physical_cond_columns(data, particle_conditioning, material_conditioning),
-        ]
-    ).astype(np.float32)
-
-    # In "physical" mode cond_cat's first two columns are only a
-    # reporting/router convenience — ConditionEncoder never reads them
-    # (giant/model/network.py) — so a species/material outside the training
-    # vocab (the whole point of physical-property conditioning) gets a dummy
-    # index instead of raising. In "embedding" mode those columns ARE the
-    # conditioning signal, so an unmapped value must still raise loudly
-    # rather than silently misassign. In "onehot" mode they again go unread
-    # (the topN columns below are the real signal), so they're as permissive
-    # as "physical". Each axis's strictness is independent.
-    pdg_strict = particle_conditioning == "embedding"
-    mat_strict = material_conditioning == "embedding"
-    pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=pdg_strict)
-    mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=mat_strict)
-    cat_cols = [pdg_idx, mat_idx]
-    if pdg_topn_map is not None:
-        cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map))
-    if mat_topn_map is not None:
-        cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map))
-    cond_cat = np.column_stack(cat_cols)
+    layout = CondLayout.from_types(particle_conditioning, material_conditioning)
+    cond_cont, cond_cat = _build_cond_arrays(data, pdg_map, mat_map, layout, pdg_topn_map, mat_topn_map)
 
     if cond_normalizer is not None:
-        cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, particle_conditioning, material_conditioning)
+        cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, layout)
 
     return cond_cont, cond_cat
 
@@ -823,8 +853,7 @@ def build_cond_features(
 def _cond_normalizer_transform(
     cond_cont: np.ndarray,
     cond_normalizer: "Normalizer",
-    particle_conditioning: str,
-    material_conditioning: str,
+    layout: CondLayout,
 ) -> np.ndarray:
     """Apply ``cond_normalizer``, padding a legacy narrower normalizer if needed.
 
@@ -832,7 +861,7 @@ def _cond_normalizer_transform(
     8->15, ``giant/constants.py``) saved a ``COND_DIM_BASE``-wide (8) cond
     normalizer, fit before ``build_cond_features`` grew the extra physical
     columns. When NEITHER axis is "physical" those columns are never read by
-    ``ConditionEncoder`` (``giant/model/network.py``), so padding the missing
+    ``ConditionEncoder`` (``giant/model/encoders.py``), so padding the missing
     entries with mean=0/std=1 is a safe no-op that keeps such checkpoints
     usable under the current, always-``COND_DIM``-wide contract. If EITHER
     axis is "physical" its columns are load-bearing, so a mismatch there is a
@@ -843,14 +872,14 @@ def _cond_normalizer_transform(
     width = cond_cont.shape[-1]
     if mean.shape[-1] < width:
         physical_load_bearing = "physical" in (
-            particle_conditioning,
-            material_conditioning,
+            layout.particle_type,
+            layout.material_type,
         )
         if physical_load_bearing:
             raise ValueError(
                 f"cond normalizer has {mean.shape[-1]} columns, expected "
-                f"{width}, and particle_conditioning={particle_conditioning!r}/"
-                f"material_conditioning={material_conditioning!r} reads the "
+                f"{width}, and particle_conditioning={layout.particle_type!r}/"
+                f"material_conditioning={layout.material_type!r} reads the "
                 "physical columns directly — this checkpoint predates "
                 "physical-property conditioning and can't be safely padded; "
                 "retrain it under the current code."
@@ -928,7 +957,7 @@ def build_features(
         instead) for callers (normalizer fitting) that only read
         `sec_cont[:, :, 4:6]` and would otherwise discard that work.
 
-    pdg_topn_map/mat_topn_map: appended `cond_cat` columns for
+    pdg_topn_map/mat_topn_map: source of the extra `cond_cat` columns for
         `ConditionEncoder`'s `"onehot"` mode — see `build_cond_features`.
 
     sec_type_class_map: the map `sec_type_idx` is looked up against — a
@@ -959,29 +988,8 @@ def build_features(
     ).astype(np.float32)  # (N, 9)
 
     # Phase 2: conditioning drops n_sec and log(e_sec)
-    cond_cont = np.column_stack(
-        [
-            data["pre_pos"],
-            log_transform(data["pre_E"]),
-            data["pre_dir"],
-            data["layer_id"].astype(np.float32),
-        ]
-    ).astype(np.float32)  # (N, COND_DIM_BASE=8)
-    cond_cont = np.column_stack(
-        [
-            cond_cont,
-            _physical_cond_columns(data, particle_conditioning, material_conditioning),
-        ]
-    ).astype(np.float32)  # (N, COND_DIM=15)
-
-    pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map)
-    mat_idx = _vectorized_map_lookup(data["material"], mat_map)
-    cat_cols = [pdg_idx, mat_idx]
-    if pdg_topn_map is not None:
-        cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map))
-    if mat_topn_map is not None:
-        cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map))
-    cond_cat = np.column_stack(cat_cols)  # (N, 2/3/4)
+    layout = CondLayout.from_types(particle_conditioning, material_conditioning)
+    cond_cont, cond_cat = _build_cond_arrays(data, pdg_map, mat_map, layout, pdg_topn_map, mat_topn_map)
 
     n_sec_raw = data["n_sec"].astype(np.int64)  # (N,) unclamped, for the valid-slot mask
 
@@ -1048,7 +1056,7 @@ def build_features(
         target_normalizer = Normalizer().fit(target_s1)
 
     if cond_normalizer is not None:
-        cond_cont = cond_normalizer.transform(cond_cont)
+        cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, layout)
     if target_normalizer is not None:
         target_s1 = target_normalizer.transform(target_s1)
     if sec_phys_normalizer is not None:

giant/model/builders.py

@@ -15,6 +15,7 @@ from giant.model.models import (
     resolve_type_n_classes,
     stage2_trunk_sec_dim,
 )
+from giant.model.objectives import build_objective
 from giant.model.routers import Router, _build_router_from_cfg
 
 # ---------------------------------------------------------------------------
@@ -44,11 +45,10 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
     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)
+    conditioning_cfg = ConditioningConfig.from_dict(cfg["conditioning"])
+    particle_cfg = conditioning_cfg.particle
+    material_cfg = conditioning_cfg.material
+    particle_conditioning = particle_cfg.type
     s1_spec = Stage1ModelConfig.from_dict(cfg["stage1_model"])
     s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"])
     cond_out_dim = conditioning_cfg.out_dim
@@ -64,10 +64,11 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
         if s1_spec.router.enabled:
             stage1_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning)
         generator = s1_spec.generator
+        objective = build_objective(generator)
         # wgan has no time_dim concept (no diffusion/flow time variable) —
         # matches the pre-dataclass .get("time_dim", 64) fallback, which
         # always hit its default for a wgan sub-block too.
-        time_dim = getattr(s1_spec, generator).time_dim if generator != "wgan" else 64
+        time_dim = getattr(s1_spec, generator).time_dim if objective.needs_time else 64
         n_sec_owner = s2_spec.n_sec.owner
         n_sec_head_k_max = s2_spec.k_max if n_sec_owner == "stage1" else None
         result["stage1"] = Stage1Model(
@@ -83,8 +84,11 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
             time_dim=time_dim,
             noise_dim=s1_spec.wgan.noise_dim,
             router=stage1_router,
+            trunk_type=s1_spec.trunk.type,
+            block_conditioning=s1_spec.trunk.block_conditioning,
             n_sec_head_k_max=n_sec_head_k_max,
             cond_enc=shared_cond_enc,
+            n_sec_head_cfg=s1_spec.heads.n_sec.to_dict(),
         )
 
     if s2_spec.active:
@@ -97,12 +101,14 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
             else:
                 stage2_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning)
         generator = s2_spec.generator
+        objective = build_objective(generator)
         # wgan has no time_dim concept — see the matching comment in stage 1
         # above.
-        time_dim = getattr(s2_spec, generator).time_dim if generator != "wgan" else 64
+        time_dim = getattr(s2_spec, generator).time_dim if objective.needs_time else 64
         n_sec_owner = s2_spec.n_sec.owner
+        stop_token = s2_spec.n_sec.mode == "stop_token"
         k_max = s2_spec.k_max
-        particle_type_cfg = s2_spec.particle_type.to_dict()
+        particle_type_cfg = s2_spec.particle_type
 
         if decoder == "autoregressive":
             ar_cfg = s2_spec.autoregressive
@@ -121,16 +127,23 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 noise_dim=s2_spec.wgan.noise_dim,
                 k_max=k_max,
                 router=stage2_router,
-                build_n_sec_head=n_sec_owner != "stage1",
+                trunk_type=s2_spec.trunk.type,
+                block_conditioning=s2_spec.trunk.block_conditioning,
+                build_n_sec_head=n_sec_owner != "stage1" and not stop_token,
                 particle_type_cfg=particle_type_cfg,
                 history=ar_cfg.history,
                 attn_n_heads=ar_cfg.attn_n_heads,
                 attn_n_layers=ar_cfg.attn_n_layers,
                 cond_enc=shared_cond_enc,
+                n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
+                type_head_cfg=s2_spec.heads.type.to_dict(),
+                build_stop_head=stop_token,
+                stop_sampling=s2_spec.n_sec.stop_sampling,
+                stop_head_cfg=s2_spec.heads.n_sec.to_dict(),
             )
         else:
             sec_dim = stage2_trunk_sec_dim(
-                particle_type_cfg, generator, k_max, resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
+                particle_type_cfg, generator, k_max, resolve_type_n_classes(particle_type_cfg, particle_cfg.emb_dim)
             )
             result["stage2"] = Stage2OneShot(
                 pdg_vocab=pdg_vocab,
@@ -148,9 +161,13 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 noise_dim=s2_spec.wgan.noise_dim,
                 k_max=k_max,
                 router=stage2_router,
+                trunk_type=s2_spec.trunk.type,
+                block_conditioning=s2_spec.trunk.block_conditioning,
                 build_n_sec_head=n_sec_owner != "stage1",
                 particle_type_cfg=particle_type_cfg,
                 cond_enc=shared_cond_enc,
+                n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
+                type_head_cfg=s2_spec.heads.type.to_dict(),
             )
 
     return result
@@ -164,17 +181,16 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
     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)
+    conditioning_cfg = ConditioningConfig.from_dict(cfg["conditioning"])
+    particle_cfg = conditioning_cfg.particle
+    material_cfg = conditioning_cfg.material
     cond_out_dim = conditioning_cfg.out_dim
     s1_spec = Stage1ModelConfig.from_dict(cfg["stage1_model"])
     s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"])
 
     result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None}
 
-    if s1_spec.active and s1_spec.generator == "wgan":
+    if s1_spec.active and build_objective(s1_spec.generator).is_adversarial:
         result["stage1"] = CriticModel(
             pdg_vocab=pdg_vocab,
             mat_vocab=mat_vocab,
@@ -188,11 +204,14 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
             stage="stage1",
         )
 
-    if s2_spec.active and s2_spec.generator == "wgan":
+    if s2_spec.active and build_objective(s2_spec.generator).is_adversarial:
         k_max = s2_spec.k_max
-        particle_type_cfg = s2_spec.particle_type.to_dict()
+        particle_type_cfg = s2_spec.particle_type
         in_dim = stage2_trunk_sec_dim(
-            particle_type_cfg, "wgan", k_max, resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
+            particle_type_cfg,
+            s2_spec.generator,
+            k_max,
+            resolve_type_n_classes(particle_type_cfg, particle_cfg.emb_dim),
         )
         result["stage2"] = CriticModel(
             pdg_vocab=pdg_vocab,

giant/model/encoders.py

@@ -5,82 +5,63 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from giant.cond_layout import CondLayout
+from giant.config import ConditioningAxisConfig
 from giant.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM
 from giant.model.layers import _make_axis_mlp
 
 
-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:
+    (`particle_cfg`/`material_cfg`, each a `ConditioningAxisConfig`) and may
+    mix freely, e.g. material "physical" with particle "embedding". Three
+    modes per axis:
       - "embedding": a learned `nn.Embedding` lookup, indexed by `cond_cat`'s
         dense training-vocab index. Memorizes the training menu.
       - "physical": an `n_layers`-deep MLP over the axis's raw physical
-        properties (already present in `cond_cont[:, COND_DIM_BASE:]` — see
+        properties (already present in `cond_cont`'s physical block — see
         giant.data.transforms.build_features), computable for any PDG code /
         material name rather than only ones seen in training.
       - "onehot": a fixed, unlearned one-hot vector over a top-N-plus-other
         class map (`giant.data.loader.build_topn_map_from_files`/
         `build_pdg_topn_map_from_files`), read from `cond_cat`'s extra
-        top-N-index column(s) — see `_cat_col_layout`.
+        top-N-index column(s).
+
+    Every column index/slice comes from `self.layout`
+    (`giant.cond_layout.CondLayout`), the same object the feature builders
+    lay the arrays out with, so the two sides cannot drift apart.
     """
 
     def __init__(
         self,
         pdg_vocab: int,
         mat_vocab: int,
-        particle_cfg: dict,
-        material_cfg: dict,
+        particle_cfg: ConditioningAxisConfig,
+        material_cfg: ConditioningAxisConfig,
         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"])
+        self.particle_cfg = particle_cfg
+        self.material_cfg = material_cfg
+        # Also validates both axis types — an unknown one raises here.
+        self.layout = CondLayout.from_types(particle_cfg.type, material_cfg.type)
 
-        p_type = particle_cfg["type"]
-        p_emb_dim = particle_cfg["emb_dim"]
+        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}")
+            self.particle_mlp = _make_axis_mlp(PARTICLE_PHYS_DIM, p_emb_dim, particle_cfg.n_layers)
 
-        m_type = material_cfg["type"]
-        m_emb_dim = material_cfg["emb_dim"]
+        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}")
+            self.material_mlp = _make_axis_mlp(MATERIAL_PHYS_DIM, m_emb_dim, material_cfg.n_layers)
 
         in_dim = COND_DIM_BASE + p_emb_dim + m_emb_dim
         self.mlp = nn.Sequential(
@@ -90,33 +71,31 @@ class ConditionEncoder(nn.Module):
         )
 
     def _particle_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor):
-        p_type = self.particle_cfg["type"]
+        p_type = self.particle_cfg.type
         if p_type == "embedding":
-            return self.pdg_emb(cond_cat[:, 0])
+            return self.pdg_emb(cond_cat[:, self.layout.PDG_COL])
         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 self.particle_mlp(cond_cont[:, self.layout.particle_phys])
+        assert self.layout.particle_topn_col is not None
         return F.one_hot(
-            cond_cat[:, self._particle_topn_col],
-            num_classes=self.particle_cfg["emb_dim"],
+            cond_cat[:, self.layout.particle_topn_col],
+            num_classes=self.particle_cfg.emb_dim,
         ).float()
 
     def _material_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor):
-        m_type = self.material_cfg["type"]
+        m_type = self.material_cfg.type
         if m_type == "embedding":
-            return self.mat_emb(cond_cat[:, 1])
+            return self.mat_emb(cond_cat[:, self.layout.MAT_COL])
         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 self.material_mlp(cond_cont[:, self.layout.material_phys])
+        assert self.layout.material_topn_col is not None
         return F.one_hot(
-            cond_cat[:, self._material_topn_col],
-            num_classes=self.material_cfg["emb_dim"],
+            cond_cat[:, self.layout.material_topn_col],
+            num_classes=self.material_cfg.emb_dim,
         ).float()
 
     def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         pdg_e = self._particle_embed(cond_cont, cond_cat)
         mat_e = self._material_embed(cond_cont, cond_cat)
-        x = torch.cat([cond_cont[:, :COND_DIM_BASE], pdg_e, mat_e], dim=-1)
+        x = torch.cat([cond_cont[:, self.layout.base], pdg_e, mat_e], dim=-1)
         return self.mlp(x)

giant/model/history.py

@@ -1,5 +1,9 @@
 """History encoders — stage-2 autoregressive only. Self-contained, no
-dependency on any other `giant.model` submodule (issues.md Issue 8)."""
+dependency on any other `giant.model` submodule (issues.md Issue 8), except
+for the `HISTORY_REGISTRY`/`build_history` factory, which mirrors
+`giant.model.routers`'s `Router`/`ROUTER_REGISTRY` pattern (gitea #35)."""
+
+import inspect
 
 import torch
 import torch.nn as nn
@@ -9,19 +13,57 @@ 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)."""
+    and `AttentionHistory` are the two registered implementations (see
+    `HISTORY_REGISTRY`/`build_history`). Inference (`giant/sample.py`)
+    generates one token at a time and cannot afford `forward`'s per-step cost
+    to be O(K) (attention would then be O(K^2) over a rollout's k_max loop),
+    so this interface also declares `init_cache`/`step` for that incremental
+    path, with working O(1) defaults here (`init_cache` -> `None`, `step` ->
+    one `forward` call ignoring `cache`) — correct for any encoder whose
+    per-step cost is already O(1) (i.e. it only ever looks at the previous
+    token, not the full prefix), which is what `MarkovHistory` relies on.
+    `AttentionHistory` overrides both with real incremental-cache versions,
+    since its `forward` genuinely needs the full prefix."""
 
     def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
         raise NotImplementedError
 
+    def init_cache(self) -> object:
+        return None
 
+    def step(self, feat: torch.Tensor, has_prev: torch.Tensor, cache: object) -> tuple[torch.Tensor, object]:
+        return self.forward(feat, has_prev), cache
+
+
+HISTORY_REGISTRY: dict[str, type[HistoryEncoder]] = {}
+
+
+def register_history(name: str):
+    def decorator(cls: type[HistoryEncoder]) -> type[HistoryEncoder]:
+        HISTORY_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def build_history(name: str, in_dim: int, out_dim: int, **kwargs) -> HistoryEncoder:
+    """Factory: look up a `HistoryEncoder` subclass by name from the registry.
+
+    Every registered history type is fed the same `stage2_model.autoregressive`
+    kwargs; kwargs not declared by that type's constructor are silently
+    dropped, so per-type hyperparameters (e.g. `AttentionHistory`'s
+    `n_heads`/`n_layers`) can coexist in one config without special-casing —
+    mirrors `giant.model.routers.build_router`.
+    """
+    if name not in HISTORY_REGISTRY:
+        raise ValueError(f"unknown history type {name!r}; available: {sorted(HISTORY_REGISTRY)}")
+    cls = HISTORY_REGISTRY[name]
+    accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "in_dim", "out_dim"}
+    filtered = {k: v for k, v in kwargs.items() if k in accepted}
+    return cls(in_dim, out_dim, **filtered)
+
+
+@register_history("markov")
 class MarkovHistory(HistoryEncoder):
     """Summarizes the previous secondary's own `(energy_fraction, direction,
     type_representation)` through one small MLP — the "markov" history:
@@ -90,6 +132,7 @@ class _CausalAttnBlock(nn.Module):
         return x, kv
 
 
+@register_history("attention")
 class AttentionHistory(HistoryEncoder):
     """Causal self-attention over the emitted-token prefix — the more
     expressive alternative to `MarkovHistory`'s fixed previous-token-only
@@ -137,16 +180,19 @@ class AttentionHistory(HistoryEncoder):
 
     def step(
         self,
-        token_feat: torch.Tensor,
+        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`).
+        cache: object,
+    ) -> tuple[torch.Tensor, object]:
+        """`feat`/`has_prev`: `(B, 1, in_dim)`/`(B, 1)` — the newest token's
+        own features (what would be `feat[:, k]` in `forward`). `cache`: the
+        `list[Tensor | None]` from `init_cache`/a previous `step` call (typed
+        `object` here to match `HistoryEncoder.step`'s base signature).
         Advances every block's cache by this position and returns this
         position's output (`(B, 1, out_dim)`, the correct history summary for
         the NEXT slot) plus the updated cache."""
-        x = self._embed(token_feat, has_prev)
+        assert isinstance(cache, list)
+        x = self._embed(feat, has_prev)
         new_cache: list[torch.Tensor | None] = []
         for block, kv in zip(self.blocks, cache):
             x, kv_new = block.step(x, kv)

giant/model/layers.py

@@ -42,6 +42,31 @@ def _make_axis_mlp(in_dim: int, emb_dim: int, n_layers: int) -> nn.Sequential:
     return nn.Sequential(*layers)
 
 
+def build_mlp_head(
+    in_dim: int, out_dim: int, hidden: int, depth: int = 2, act: type[nn.Module] = nn.SiLU
+) -> nn.Sequential:
+    """`depth`-layer MLP head (gitea #36) — factors out the n_sec_head/
+    type_head pattern duplicated five times across `giant.model.models`.
+
+    `depth=1` is a bare `Linear(in_dim, out_dim)` (no hidden layer/
+    activation); `depth>=2` is `Linear(in_dim, hidden) -> act -> [Linear
+    (hidden, hidden) -> act] * (depth-2) -> Linear(hidden, out_dim)` —
+    `depth=2` reproduces every pre-#36 n_sec_head/type_head exactly when
+    `hidden == hidden_dim // 2`. Mirrors `_make_axis_mlp`'s depth
+    convention above, but takes `hidden` and `out_dim` as independent
+    widths (n_sec_head/type_head's hidden width is not their output width,
+    unlike the particle/material axis MLPs)."""
+    if depth < 1:
+        raise ValueError(f"depth must be >= 1, got {depth}")
+    if depth == 1:
+        return nn.Sequential(nn.Linear(in_dim, out_dim))
+    layers: list[nn.Module] = [nn.Linear(in_dim, hidden), act()]
+    for _ in range(depth - 2):
+        layers += [nn.Linear(hidden, hidden), act()]
+    layers.append(nn.Linear(hidden, out_dim))
+    return nn.Sequential(*layers)
+
+
 class ContextAdapter(nn.Module):
     """Projects a stage's outcome (e.g. Stage 1's 9D target) down to a
     fixed-width context vector for a downstream stage's conditioning —
@@ -57,6 +82,26 @@ class ContextAdapter(nn.Module):
         return torch.tanh(self.proj(x))
 
 
+BLOCK_REGISTRY: dict[str, type[nn.Module]] = {}
+
+
+def register_block(name: str):
+    def decorator(cls: type[nn.Module]) -> type[nn.Module]:
+        BLOCK_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def build_block(name: str, dim: int, cond_dim: int, dropout: float = 0.0) -> nn.Module:
+    """Factory: look up a registered conditioning-injection block by name and
+    construct one instance — `trunk.block_conditioning` (gitea #34)."""
+    if name not in BLOCK_REGISTRY:
+        raise ValueError(f"unknown block conditioning type {name!r}; available: {sorted(BLOCK_REGISTRY)}")
+    return BLOCK_REGISTRY[name](dim, cond_dim, dropout)
+
+
+@register_block("add")
 class ResBlock(nn.Module):
     def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None:
         super().__init__()
@@ -74,3 +119,66 @@ class ResBlock(nn.Module):
         h = self.dropout(h)
         h = self.linear2(h)
         return x + h
+
+
+@register_block("film")
+class FilmResBlock(nn.Module):
+    """FiLM conditioning (Perez et al. 2018): a per-channel scale+shift
+    modulates the normalized features, on top of the norm's own affine —
+    an *additional* modulation, unlike `AdaLNResBlock` below, which replaces
+    the norm's affine outright. `film_proj` is zero-initialized so
+    `gamma=beta=0` at construction — conditioning has no effect on the
+    output until training moves it, a stable starting point (though not a
+    literal identity block, since `linear1`/`linear2` aren't zero-init)."""
+
+    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.linear1 = nn.Linear(dim, dim)
+        self.film_proj = nn.Linear(cond_dim, 2 * dim)
+        nn.init.zeros_(self.film_proj.weight)
+        nn.init.zeros_(self.film_proj.bias)
+        self.act = nn.SiLU()
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        h = self.norm(x)
+        gamma, beta = self.film_proj(cond).chunk(2, dim=-1)
+        h = h * (1 + gamma) + beta
+        h = self.linear1(h)
+        h = self.act(h)
+        h = self.dropout(h)
+        h = self.linear2(h)
+        return x + h
+
+
+@register_block("adaln")
+class AdaLNResBlock(nn.Module):
+    """AdaLN-Zero conditioning (DiT, Peebles & Xie 2022): the norm's own
+    affine is replaced by a conditioning-derived scale/shift, and the
+    residual branch is scaled by a conditioning-derived gate. `adaln_proj`
+    is zero-initialized, so `scale=shift=gate=0` at construction — the block
+    is the exact identity function at init (`x + 0 * h' == x`), regardless
+    of `x`/`cond`."""
+
+    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False)
+        self.linear1 = nn.Linear(dim, dim)
+        self.adaln_proj = nn.Linear(cond_dim, 3 * dim)
+        nn.init.zeros_(self.adaln_proj.weight)
+        nn.init.zeros_(self.adaln_proj.bias)
+        self.act = nn.SiLU()
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        h = self.norm(x)
+        scale, shift, gate = self.adaln_proj(cond).chunk(3, dim=-1)
+        h = h * (1 + scale) + shift
+        h = self.linear1(h)
+        h = self.act(h)
+        h = self.dropout(h)
+        h = self.linear2(h)
+        return x + gate * h

giant/model/models.py

@@ -4,10 +4,12 @@
 import torch
 import torch.nn as nn
 
+from giant.config import ConditioningAxisConfig, HeadConfig, ParticleTypeConfig
 from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, SEC_SLOT_DIM, X_DIM
 from giant.model.encoders import ConditionEncoder
-from giant.model.history import AttentionHistory, HistoryEncoder, MarkovHistory
-from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding
+from giant.model.history import HistoryEncoder, build_history
+from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding, build_mlp_head
+from giant.model.objectives import build_objective
 from giant.model.routers import Router
 from giant.model.trunks import build_trunk
 
@@ -16,7 +18,7 @@ from giant.model.trunks import build_trunk
 # ---------------------------------------------------------------------------
 
 
-def resolve_type_n_classes(particle_type_cfg: dict, particle_emb_dim: int) -> int:
+def resolve_type_n_classes(particle_type_cfg: ParticleTypeConfig, particle_emb_dim: int) -> int:
     """Effective width fed to `stage2_type_dim`/`stage2_trunk_sec_dim` in
     place of a bare `conditioning.particle.emb_dim` read. Under
     `target = "onehot"` this is `stage2_model.particle_type.n_classes` (0 =
@@ -27,50 +29,52 @@ def resolve_type_n_classes(particle_type_cfg: dict, particle_emb_dim: int) -> in
     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
+    if particle_type_cfg.target == "onehot":
+        return particle_type_cfg.n_classes or particle_emb_dim
     return particle_emb_dim
 
 
-def stage2_type_dim(particle_type_cfg: dict, emb_dim: int) -> int:
+def stage2_type_dim(particle_type_cfg: ParticleTypeConfig, emb_dim: int) -> int:
     """Width of a single secondary slot's type slice —
     `PARTICLE_PHYS_DIM` (log_mass, charge) for `target = "physical"`, else
     `emb_dim` (both `"onehot"` class logits and `"embedding"` vectors are
     this many classes/dims wide — callers resolve `emb_dim` via
     `resolve_type_n_classes` first)."""
-    target = particle_type_cfg.get("target", "physical")
-    return PARTICLE_PHYS_DIM if target == "physical" else emb_dim
+    return PARTICLE_PHYS_DIM if particle_type_cfg.target == "physical" else emb_dim
 
 
-def stage2_trunk_sec_dim(particle_type_cfg: dict, generator: str, k_max: int, emb_dim: int) -> int:
+def stage2_trunk_sec_dim(particle_type_cfg: ParticleTypeConfig, generator: str, k_max: int, emb_dim: int) -> int:
     """`Stage2OneShot`'s trunk output width.
 
     `target = "physical"` is untouched from v0.2/today:
     `k_max * SEC_SLOT_DIM`, the type slice folded into the same
     flow-matched/WGAN vector as the continuous stick/dir slots.
 
-    `target` in `("onehot", "embedding")`: under `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` in `("onehot", "embedding")`: under an objective with
+    `folds_type_slice` (currently just wgan) the type slice is still folded
+    in (adversarial for onehot via ST-Gumbel, already-continuous for
+    embedding), just `emb_dim` wide instead of `PARTICLE_PHYS_DIM` wide:
+    `k_max * (CONT_SLOT_DIM + emb_dim)`. Otherwise (flow/ddpm) the type slice
+    isn't part of this vector at all — it's `Stage2OneShot.type_head`'s job
+    instead — so the trunk only covers `k_max * CONT_SLOT_DIM`.
     """
-    target = particle_type_cfg.get("target", "physical")
-    if target == "physical":
+    if particle_type_cfg.target == "physical":
         return k_max * SEC_SLOT_DIM
-    if generator == "wgan":
+    if build_objective(generator).folds_type_slice:
         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`).
+class StageModel(nn.Module):
+    """Base owning the scaffolding common to `Stage1Model`, `Stage2OneShot`,
+    `Stage2Autoregressive` (gitea #39): build-or-share `cond_enc`,
+    `particle_type_cfg` normalisation, and — via `_build_trunk_and_heads`,
+    called by each subclass's `__init__` once its own conditioning-assembly
+    modules exist — the objective/time-embedding/trunk construction and the
+    `n_sec_head`/`type_head` classifier heads. A subclass supplies only its
+    own conditioning assembly (`Stage1Model` uses `cond_enc` directly;
+    `Stage2OneShot`/`Stage2Autoregressive` add a context-fusion path) and its
+    trunk's output width.
 
     `cond_enc`, if given, is used in place of building a fresh
     `ConditionEncoder` — `conditioning.share_stages = true`: `build_models`
@@ -81,8 +85,149 @@ class Stage1Model(nn.Module):
         self,
         pdg_vocab: int,
         mat_vocab: int,
-        particle_cfg: dict,
-        material_cfg: dict,
+        particle_cfg: ConditioningAxisConfig,
+        material_cfg: ConditioningAxisConfig,
+        cond_out_dim: int,
+        generator: str,
+        noise_dim: int,
+        k_max: int | None = None,
+        particle_type_cfg: ParticleTypeConfig | None = None,
+        cond_enc: ConditionEncoder | None = None,
+    ) -> None:
+        super().__init__()
+        self.generator_kind = generator
+        self.noise_dim = noise_dim
+        self.k_max = k_max
+        # `ParticleTypeConfig()`'s own dataclass default is target="onehot"
+        # (the config.toml default when [stage2_model.particle_type] is
+        # omitted) — a different question from "nobody passed anything to
+        # this constructor", which direct/test construction relies on
+        # defaulting to "physical" (build_models/build_critics always pass
+        # particle_type_cfg explicitly, so this sentinel is never hit there).
+        self.particle_type_cfg = (
+            particle_type_cfg if particle_type_cfg is not None else ParticleTypeConfig(target="physical")
+        )
+        self.type_dim = stage2_type_dim(
+            self.particle_type_cfg, resolve_type_n_classes(self.particle_type_cfg, particle_cfg.emb_dim)
+        )
+        self.cond_enc = (
+            cond_enc
+            if cond_enc is not None
+            else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
+        )
+
+    def _build_trunk_and_heads(
+        self,
+        *,
+        trunk_out_dim: int,
+        hidden_dim: int,
+        n_res_blocks: int,
+        cond_out_dim: int,
+        time_dim: int,
+        router: Router | None,
+        trunk_type: str,
+        block_conditioning: str,
+        dropout: float,
+        n_sec_head_k_max: int | None,
+        n_sec_head_cfg: dict | None,
+        type_head_out_dim: int | None,
+        type_head_cfg: dict | None,
+        build_stop_head: bool = False,
+        stop_head_cfg: dict | None = None,
+    ) -> None:
+        """Builds `self.time_emb`, `self.trunk`, `self.n_sec_head`,
+        `self.type_head`, `self.stop_head`. Called by a subclass's `__init__`
+        after it has set up its own conditioning-assembly modules —
+        `merged_cond_dim` below must match the width that assembly
+        (`_cond_embed`/`_base_cond`/`_token_cond`, or plain `cond_enc` for
+        `Stage1Model`) actually produces.
+
+        `n_sec_head` is built iff `n_sec_head_k_max is not None` (output
+        width `n_sec_head_k_max + 1`) — `Stage1Model` passes this only for a
+        migrated v0.2 checkpoint, `Stage2OneShot`/`Stage2Autoregressive` pass
+        it whenever `build_n_sec_head=True`. `type_head` is built iff
+        `type_head_out_dim is not None` (the caller — only the two Stage2
+        classes — passes `None` exactly when `particle_type_cfg.target ==
+        "physical"`) *and* the objective doesn't fold the type slice into its
+        own trunk output (checked here, since `objective` is already needed
+        for the trunk itself). `stop_head` is built iff `build_stop_head` —
+        only `Stage2Autoregressive` ever passes `True` (`n_sec.mode ==
+        "stop_token"`, mutually exclusive with `n_sec_head`), a single
+        `cond_out_dim -> 1` logit per call, same `HeadConfig` shape rules as
+        the other two heads.
+        """
+        objective = build_objective(self.generator_kind)
+        has_time = objective.needs_time
+        self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
+        merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
+        in_dim = objective.trunk_in_dim(trunk_out_dim, self.noise_dim)
+        self.trunk = build_trunk(
+            router,
+            trunk_type,
+            in_dim,
+            trunk_out_dim,
+            hidden_dim,
+            n_res_blocks,
+            merged_cond_dim,
+            dropout,
+            block_conditioning,
+        )
+        self.n_sec_head = None
+        if n_sec_head_k_max is not None:
+            head_cfg = HeadConfig.from_dict(n_sec_head_cfg)
+            hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
+            self.n_sec_head = build_mlp_head(cond_out_dim, n_sec_head_k_max + 1, hidden, head_cfg.depth)
+        self.type_head = None
+        if type_head_out_dim is not None and not objective.folds_type_slice:
+            head_cfg = HeadConfig.from_dict(type_head_cfg)
+            hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
+            self.type_head = build_mlp_head(cond_out_dim, type_head_out_dim, hidden, head_cfg.depth)
+        self.stop_head = None
+        if build_stop_head:
+            head_cfg = HeadConfig.from_dict(stop_head_cfg)
+            hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
+            self.stop_head = build_mlp_head(cond_out_dim, 1, hidden, head_cfg.depth)
+
+    def _require_n_sec_head(self) -> None:
+        if self.n_sec_head is None:
+            raise RuntimeError(
+                f"this {type(self).__name__} has no n_sec_head — it belongs to "
+                "a migrated v0.2 checkpoint (n_sec.owner='stage1'); call "
+                "stage1.predict_n_sec(cond_cont, cond_cat) instead"
+            )
+
+    def _require_type_head(self) -> None:
+        if self.type_head is None:
+            raise RuntimeError(
+                f"this {type(self).__name__} has no type_head — either "
+                "particle_type.target='physical' (the type slice is part of "
+                "forward()'s own output) or generator='wgan' (the WGAN "
+                "trainer reads the type slice out of forward()'s output "
+                "directly instead)"
+            )
+
+    def _require_stop_head(self) -> None:
+        if self.stop_head is None:
+            raise RuntimeError(
+                f"this {type(self).__name__} has no stop_head — only a "
+                "Stage2Autoregressive built with stage2_model.n_sec.mode = "
+                "'stop_token' owns one"
+            )
+
+
+class Stage1Model(StageModel):
+    """Predicts the 9D primary post-step vector. No `n_sec_head` — fresh runs
+    move it to stage 2, except for a migrated v0.2 checkpoint
+    (`n_sec_head_k_max` given), where it stays attached here
+    since that's where its weights live and what conditioning it was trained
+    against (see `_migrate_legacy_model_config`)."""
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        particle_cfg: ConditioningAxisConfig,
+        material_cfg: ConditioningAxisConfig,
         hidden_dim: int = 256,
         n_res_blocks: int = 6,
         cond_out_dim: int = 128,
@@ -92,29 +237,37 @@ class Stage1Model(nn.Module):
         time_dim: int = 64,
         noise_dim: int = 64,
         router: Router | None = None,
+        trunk_type: str = "resmlp",
+        block_conditioning: str = "add",
         n_sec_head_k_max: int | None = None,
         cond_enc: ConditionEncoder | None = None,
+        n_sec_head_cfg: dict | None = None,
     ) -> None:
-        super().__init__()
-        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)
+        super().__init__(
+            pdg_vocab,
+            mat_vocab,
+            particle_cfg,
+            material_cfg,
+            cond_out_dim=cond_out_dim,
+            generator=generator,
+            noise_dim=noise_dim,
+            cond_enc=cond_enc,
+        )
+        self._build_trunk_and_heads(
+            trunk_out_dim=x_dim,
+            hidden_dim=hidden_dim,
+            n_res_blocks=n_res_blocks,
+            cond_out_dim=cond_out_dim,
+            time_dim=time_dim,
+            router=router,
+            trunk_type=trunk_type,
+            block_conditioning=block_conditioning,
+            dropout=dropout,
+            n_sec_head_k_max=n_sec_head_k_max,
+            n_sec_head_cfg=n_sec_head_cfg,
+            type_head_out_dim=None,
+            type_head_cfg=None,
         )
-        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,
@@ -127,21 +280,28 @@ class Stage1Model(nn.Module):
         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."""
+    def _require_n_sec_head(self) -> None:
+        """Overrides `StageModel`'s guard — a `Stage1Model` with no
+        `n_sec_head` points the caller to stage 2 (n_sec's default owner),
+        not to `stage1` as the base's message would."""
         if self.n_sec_head is None:
             raise RuntimeError(
                 "this Stage1Model has no n_sec_head — n_sec now lives on "
                 "stage 2 by default; this method only exists "
                 "for a migrated v0.2 checkpoint (n_sec.owner='stage1')"
             )
+
+    def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        """Return n_sec logits (B, K_MAX+1) from conditioning alone. Only
+        valid on a migrated v0.2 checkpoint's Stage1Model — fresh v0.3.0
+        configs predict n_sec from Stage2OneShot instead."""
+        self._require_n_sec_head()
+        assert self.n_sec_head is not None
         c_emb = self.cond_enc(cond_cont, cond_cat)
         return self.n_sec_head(c_emb)
 
 
-class Stage2OneShot(nn.Module):
+class Stage2OneShot(StageModel):
     """Predicts all `k_max` secondary slots simultaneously — v0.2 behaviour,
     reproduced exactly (`decoder = "autoregressive"` is `Stage2Autoregressive`,
     step 4/5, not implemented yet).
@@ -150,29 +310,27 @@ class Stage2OneShot(nn.Module):
     (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
+    `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
+    `"onehot"`/`"embedding"` with an objective (`giant.model.objectives`) that
+    doesn't fold the type slice (flow/ddpm), the type
     slice is predicted by a separate `type_head` instead (same shape pattern
     as `n_sec_head`) — `sec_dim` then covers only the continuous
     stick/dir slots, `type_head` covers `k_max * emb_dim` type logits/vectors.
-    Under `generator == "wgan"` the type slice stays folded into `sec_dim`
+    Under a folding objective (wgan) the type slice stays folded into `sec_dim`
     (just `emb_dim` instead of `PARTICLE_PHYS_DIM` wide) and `type_head` is
     unused (`None`) — the WGAN trainer handles the ST-Gumbel relaxation.
-
-    `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,
+        particle_cfg: ConditioningAxisConfig,
+        material_cfg: ConditioningAxisConfig,
         hidden_dim: int = 256,
         n_res_blocks: int = 6,
         cond_out_dim: int = 128,
@@ -185,51 +343,48 @@ class Stage2OneShot(nn.Module):
         noise_dim: int = 64,
         k_max: int = K_MAX,
         router: Router | None = None,
+        trunk_type: str = "resmlp",
+        block_conditioning: str = "add",
         build_n_sec_head: bool = True,
-        particle_type_cfg: dict | None = None,
+        particle_type_cfg: ParticleTypeConfig | None = None,
         cond_enc: ConditionEncoder | None = None,
+        n_sec_head_cfg: dict | None = None,
+        type_head_cfg: dict | None = None,
     ) -> None:
-        super().__init__()
-        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)
+        super().__init__(
+            pdg_vocab,
+            mat_vocab,
+            particle_cfg,
+            material_cfg,
+            cond_out_dim=cond_out_dim,
+            generator=generator,
+            noise_dim=noise_dim,
+            k_max=k_max,
+            particle_type_cfg=particle_type_cfg,
+            cond_enc=cond_enc,
         )
         self.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
+        target = self.particle_type_cfg.target
+        type_head_out_dim = None if target == "physical" else k_max * self.type_dim
+        self._build_trunk_and_heads(
+            trunk_out_dim=sec_dim,
+            hidden_dim=hidden_dim,
+            n_res_blocks=n_res_blocks,
+            cond_out_dim=cond_out_dim,
+            time_dim=time_dim,
+            router=router,
+            trunk_type=trunk_type,
+            block_conditioning=block_conditioning,
+            dropout=dropout,
+            n_sec_head_k_max=k_max if build_n_sec_head else None,
+            n_sec_head_cfg=n_sec_head_cfg,
+            type_head_out_dim=type_head_out_dim,
+            type_head_cfg=type_head_cfg,
+        )
 
     def _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)
@@ -254,12 +409,8 @@ class Stage2OneShot(nn.Module):
         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"
-            )
+        self._require_n_sec_head()
+        assert self.n_sec_head is not None
         c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
         return self.n_sec_head(c_emb)
 
@@ -273,19 +424,13 @@ class Stage2OneShot(nn.Module):
         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)"
-            )
+        self._require_type_head()
+        assert self.type_head is not None
         c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
-        return self.type_head(c_emb).view(-1, self._type_k_max, self._type_emb_dim)
+        return self.type_head(c_emb).view(-1, self.k_max, self.type_dim)
 
 
-class Stage2Autoregressive(nn.Module):
+class Stage2Autoregressive(StageModel):
     """Emits secondaries one at a time in descending-energy order, instead
     of `Stage2OneShot`'s simultaneous
     k_max-slot prediction. `history` selects `MarkovHistory` or
@@ -305,18 +450,21 @@ class Stage2Autoregressive(nn.Module):
     `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.
+    normalized slot index), and feeds `forward`/`predict_type`/`predict_stop`/
+    the trunk.
 
-    `cond_enc`, if given, is used in place of building a fresh
-    `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`).
+    `n_sec.mode = "stop_token"` (`build_stop_head=True`) replaces
+    `predict_n_sec`'s one-shot classifier with `predict_stop`'s per-token EOS
+    logit instead — the two heads are mutually exclusive (`build_n_sec_head`
+    is `False` whenever this is `True`, see `giant.model.builders`).
     """
 
     def __init__(
         self,
         pdg_vocab: int,
         mat_vocab: int,
-        particle_cfg: dict,
-        material_cfg: dict,
+        particle_cfg: ConditioningAxisConfig,
+        material_cfg: ConditioningAxisConfig,
         hidden_dim: int = 256,
         n_res_blocks: int = 6,
         cond_out_dim: int = 128,
@@ -328,29 +476,34 @@ class Stage2Autoregressive(nn.Module):
         noise_dim: int = 64,
         k_max: int = K_MAX,
         router: Router | None = None,
+        trunk_type: str = "resmlp",
+        block_conditioning: str = "add",
         build_n_sec_head: bool = True,
-        particle_type_cfg: dict | None = None,
+        particle_type_cfg: ParticleTypeConfig | None = None,
         history: str = "markov",
         attn_n_heads: int = 4,
         attn_n_layers: int = 2,
         cond_enc: ConditionEncoder | None = None,
+        n_sec_head_cfg: dict | None = None,
+        type_head_cfg: dict | None = None,
+        build_stop_head: bool = False,
+        stop_sampling: str = "greedy",
+        stop_head_cfg: dict | None = None,
     ) -> None:
-        super().__init__()
-        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)
+        super().__init__(
+            pdg_vocab,
+            mat_vocab,
+            particle_cfg,
+            material_cfg,
+            cond_out_dim=cond_out_dim,
+            generator=generator,
+            noise_dim=noise_dim,
+            k_max=k_max,
+            particle_type_cfg=particle_type_cfg,
+            cond_enc=cond_enc,
         )
+        self.history_kind = history
+        self.stop_sampling = stop_sampling
         self.context_adapter = ContextAdapter(x_dim, context_dim)
         self.base_fuse = nn.Sequential(
             nn.Linear(cond_out_dim + context_dim, cond_out_dim),
@@ -362,10 +515,8 @@ class Stage2Autoregressive(nn.Module):
         # 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)
+        self.history_encoder: HistoryEncoder = build_history(
+            history, hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers
         )
         token_fuse_in = cond_out_dim + context_dim + history_dim + 2  # +2: remaining_frac, slot_idx
         self.token_fuse = nn.Sequential(
@@ -373,37 +524,31 @@ class Stage2Autoregressive(nn.Module):
             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.type_dim` (set by StageModel.__init__) doubles as the raw
+        # `emb_dim` `stage2_trunk_sec_dim` wants: for a non-"physical" target
+        # `stage2_type_dim` already resolved `type_dim` to exactly that value;
+        # for "physical" the emb_dim argument goes unused anyway.
+        token_dim = stage2_trunk_sec_dim(self.particle_type_cfg, generator, 1, self.type_dim)
+        target = self.particle_type_cfg.target
+        type_head_out_dim = None if target == "physical" else self.type_dim
+        self._build_trunk_and_heads(
+            trunk_out_dim=token_dim,
+            hidden_dim=hidden_dim,
+            n_res_blocks=n_res_blocks,
+            cond_out_dim=cond_out_dim,
+            time_dim=time_dim,
+            router=router,
+            trunk_type=trunk_type,
+            block_conditioning=block_conditioning,
+            dropout=dropout,
+            n_sec_head_k_max=k_max if build_n_sec_head else None,
+            n_sec_head_cfg=n_sec_head_cfg,
+            type_head_out_dim=type_head_out_dim,
+            type_head_cfg=type_head_cfg,
+            build_stop_head=build_stop_head,
+            stop_head_cfg=stop_head_cfg,
         )
 
-        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)
@@ -437,12 +582,12 @@ class Stage2Autoregressive(nn.Module):
 
     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
+        (`giant/sample.py`'s AR loop) — whatever `self.history_encoder.init_cache()`
+        returns for the configured `history` type: `None` under `history="markov"`
+        (its per-step cost is already O(1) — see `HistoryEncoder`'s docstring),
+        or `AttentionHistory.init_cache()`'s real per-block KV cache under
+        `history="attention"`."""
+        return self.history_encoder.init_cache()
 
     def history_step(self, token_feat: torch.Tensor, has_prev: torch.Tensor, cache) -> tuple[torch.Tensor, object]:
         """One inference slot's worth of history encoding: advances `cache`
@@ -454,9 +599,7 @@ class Stage2Autoregressive(nn.Module):
         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
+        return self.history_encoder.step(token_feat, has_prev, cache)
 
     def forward(
         self,
@@ -496,12 +639,8 @@ class Stage2Autoregressive(nn.Module):
         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"
-            )
+        self._require_n_sec_head()
+        assert self.n_sec_head is not None
         return self.n_sec_head(self._base_cond(cond_cont, cond_cat, stage1_out))
 
     def predict_type(
@@ -515,14 +654,8 @@ class Stage2Autoregressive(nn.Module):
         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)"
-            )
+        self._require_type_head()
+        assert self.type_head is not None
         c_emb = self._token_cond(
             cond_cont,
             cond_cat,
@@ -536,6 +669,38 @@ class Stage2Autoregressive(nn.Module):
         B, K, _ = c_emb.shape
         return self.type_head(c_emb.reshape(B * K, -1)).view(B, K, self.type_dim)
 
+    def predict_stop(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+        history_feat: torch.Tensor,
+        has_prev: torch.Tensor,
+        remaining_frac: torch.Tensor,
+        slot_idx: torch.Tensor,
+        hist: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """`(B, K)` raw stop logits — `n_sec.mode = "stop_token"` only.
+        Evaluated on slot `k`'s own conditioning (which carries slot `k-1`'s
+        history, same as `predict_type`), so this is `P(n_sec == k |
+        prefix)`: a high logit at slot `k` means "stop before generating a
+        token here" — the caller (`giant.sample.sample_secondaries_ar`)
+        checks it before spending a model call on that slot's token."""
+        self._require_stop_head()
+        assert self.stop_head is not None
+        c_emb = self._token_cond(
+            cond_cont,
+            cond_cat,
+            stage1_out,
+            history_feat,
+            has_prev,
+            remaining_frac,
+            slot_idx,
+            hist=hist,
+        )
+        B, K, _ = c_emb.shape
+        return self.stop_head(c_emb.reshape(B * K, -1)).view(B, K)
+
 
 class CriticModel(nn.Module):
     """Generator-agnostic WGAN-GP critic body: a scalar realism score, for
@@ -547,8 +712,8 @@ class CriticModel(nn.Module):
         self,
         pdg_vocab: int,
         mat_vocab: int,
-        particle_cfg: dict,
-        material_cfg: dict,
+        particle_cfg: ConditioningAxisConfig,
+        material_cfg: ConditioningAxisConfig,
         in_dim: int,
         hidden_dim: int = 256,
         n_res_blocks: int = 6,

giant/model/network.py

@@ -9,18 +9,47 @@ import X` call site keeps working unchanged.
 
 from giant.model._legacy import _migrate_legacy_model_config, migrate_legacy_state_dict
 from giant.model.builders import build_critics, build_models
-from giant.model.encoders import ConditionEncoder, cat_col_layout
-from giant.model.history import AttentionHistory, HistoryEncoder, MarkovHistory, _CausalAttnBlock
-from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding, _make_axis_mlp
+from giant.model.encoders import ConditionEncoder
+from giant.model.history import (
+    HISTORY_REGISTRY,
+    AttentionHistory,
+    HistoryEncoder,
+    MarkovHistory,
+    _CausalAttnBlock,
+    build_history,
+    register_history,
+)
+from giant.model.layers import (
+    BLOCK_REGISTRY,
+    AdaLNResBlock,
+    ContextAdapter,
+    FilmResBlock,
+    ResBlock,
+    SinusoidalEmbedding,
+    _make_axis_mlp,
+    build_block,
+    build_mlp_head,
+    register_block,
+)
 from giant.model.models import (
     CriticModel,
     Stage1Model,
     Stage2Autoregressive,
     Stage2OneShot,
+    StageModel,
     resolve_type_n_classes,
     stage2_trunk_sec_dim,
     stage2_type_dim,
 )
+from giant.model.objectives import (
+    OBJECTIVE_REGISTRY,
+    DdpmObjective,
+    FlowObjective,
+    Objective,
+    WganObjective,
+    build_objective,
+    register_objective,
+)
 from giant.model.routers import (
     ROUTER_REGISTRY,
     ComposedRouter,
@@ -36,25 +65,34 @@ from giant.model.routers import (
     register_router,
 )
 from giant.model.trunks import (
+    TRUNK_REGISTRY,
     ExpertTrunk,
-    MonolithicTrunk,
     RoutedTrunk,
     Trunk,
     _route_forward,
+    build_expert_body,
     build_trunk,
+    register_trunk,
 )
 
 __all__ = [
+    "AdaLNResBlock",
     "AttentionHistory",
+    "BLOCK_REGISTRY",
     "ComposedRouter",
     "ConditionEncoder",
     "ContextAdapter",
     "CriticModel",
+    "DdpmObjective",
     "EnergyRouter",
     "ExpertTrunk",
+    "FilmResBlock",
+    "FlowObjective",
+    "HISTORY_REGISTRY",
     "HistoryEncoder",
     "MarkovHistory",
-    "MonolithicTrunk",
+    "OBJECTIVE_REGISTRY",
+    "Objective",
     "PdgRouter",
     "ProcessRouter",
     "ROUTER_REGISTRY",
@@ -65,7 +103,10 @@ __all__ = [
     "Stage1Model",
     "Stage2Autoregressive",
     "Stage2OneShot",
+    "StageModel",
+    "TRUNK_REGISTRY",
     "Trunk",
+    "WganObjective",
     "_CausalAttnBlock",
     "_build_router_from_cfg",
     "_check_router_conditioning_compat",
@@ -73,14 +114,22 @@ __all__ = [
     "_migrate_legacy_model_config",
     "_parse_composed_axes",
     "_route_forward",
+    "build_block",
     "build_composed_router",
     "build_critics",
+    "build_expert_body",
+    "build_history",
+    "build_mlp_head",
     "build_models",
+    "build_objective",
     "build_router",
     "build_trunk",
-    "cat_col_layout",
     "migrate_legacy_state_dict",
+    "register_block",
+    "register_history",
+    "register_objective",
     "register_router",
+    "register_trunk",
     "resolve_type_n_classes",
     "stage2_trunk_sec_dim",
     "stage2_type_dim",

giant/model/objectives.py

@@ -0,0 +1,202 @@
+"""Generative objectives (flow/ddpm/wgan): `Objective` base + registry,
+mirroring `giant.model.routers`'s `Router` pattern (gitea #32). Each objective
+answers, in one place, the handful of questions every stage model/sampler/
+trainer used to re-derive independently from a bare `generator` string: does
+this stage need a time embedding, is it adversarial, does it fold the
+secondary type slice into its own trunk output, what does the trunk take as
+input, which stage-1/stage-2 loss does it train against.
+
+Self-contained (no dependency on `giant.model.models`, unlike `Router` which
+`giant.model.trunks` depends on) — `Objective` never needs to construct a
+stage model or critic itself, only describe one. This also sidesteps a
+`models.py` <-> `objectives.py` import cycle, since `models.py` calls
+`build_objective`.
+"""
+
+import inspect
+
+import torch
+
+from giant.model.schedule import (
+    CosineSchedule,
+    flow_matching_loss,
+    flow_matching_loss_secondary,
+    flow_matching_loss_secondary_ar,
+)
+
+# ---------------------------------------------------------------------------
+# Objective contract
+# ---------------------------------------------------------------------------
+
+
+class Objective:
+    """Contract for a pluggable generative objective. Not an `nn.Module` —
+    unlike `Router`, no objective owns learnable parameters, so a plain
+    strategy object is the honest fit.
+
+    `needs_time`/`is_adversarial`/`folds_type_slice`/`supports_stage2_decoder`
+    are set by each concrete subclass (no defaults here — a new objective
+    should have to state all four, not silently inherit one that happens to
+    be wrong for it). See `FlowObjective`/`DdpmObjective`/`WganObjective`.
+    """
+
+    needs_time: bool
+    is_adversarial: bool
+    folds_type_slice: bool
+    supports_stage2_decoder: bool = True
+
+    def trunk_in_dim(self, out_dim: int, noise_dim: int) -> int:
+        """Width of the trunk's own input — `out_dim` (denoising/flow-matching
+        a same-shape vector) for every non-adversarial objective;
+        `WganObjective` overrides to `noise_dim` (a single-pass noise-to-output
+        generator)."""
+        return out_dim
+
+    def build_schedule(self, n_steps: int, device: torch.device) -> CosineSchedule | None:
+        """Objective-owned auxiliary state a stage trainer must build once
+        and hold onto (device-placed) across its training loop. `None` for
+        every objective except `DdpmObjective` (its noise schedule)."""
+        return None
+
+    def stage1_loss(
+        self,
+        model: torch.nn.Module,
+        x1: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        *,
+        schedule: object | None = None,
+    ) -> torch.Tensor:
+        """Stage-1 training loss. Only implemented by non-adversarial
+        objectives — `WganObjective` is unused here, `WGANStageTrainer` has
+        its own G/D step instead."""
+        raise NotImplementedError(f"{type(self).__name__} has no stage1_loss")
+
+    def stage2_loss(
+        self,
+        model: torch.nn.Module,
+        x1_s2: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_ctx: torch.Tensor,
+        sec_mask: torch.Tensor,
+        *,
+        type_dim: int | None,
+        ar_inputs: dict[str, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        """Stage-2 secondary-decoder training loss, one-shot or
+        autoregressive depending on whether `ar_inputs` is given. Same
+        adversarial caveat as `stage1_loss`."""
+        raise NotImplementedError(f"{type(self).__name__} has no stage2_loss")
+
+
+OBJECTIVE_REGISTRY: dict[str, type[Objective]] = {}
+
+
+def register_objective(name: str):
+    def decorator(cls: type[Objective]) -> type[Objective]:
+        OBJECTIVE_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def build_objective(name: str, **kwargs) -> Objective:
+    """Factory: look up an `Objective` subclass by name (a `generator`
+    config value) from the registry.
+
+    Every registered objective is fed the same kwargs; kwargs not declared by
+    that type's constructor are silently dropped, so per-type hyperparameters
+    (e.g. `DdpmObjective`'s `n_steps`) can coexist in one call without
+    special-casing — same convention as `giant.model.routers.build_router`.
+    """
+    if name not in OBJECTIVE_REGISTRY:
+        raise ValueError(f"unknown generator/objective {name!r}; available: {sorted(OBJECTIVE_REGISTRY)}")
+    cls = OBJECTIVE_REGISTRY[name]
+    accepted = set(inspect.signature(cls.__init__).parameters) - {"self"}
+    filtered = {k: v for k, v in kwargs.items() if k in accepted}
+    return cls(**filtered)
+
+
+# ---------------------------------------------------------------------------
+# Concrete objectives
+# ---------------------------------------------------------------------------
+
+
+@register_objective("flow")
+class FlowObjective(Objective):
+    """Conditional flow matching (Lipman et al. 2022) — the primary
+    objective. ~10 ODE steps at inference (`giant.sample.sample_flow`)."""
+
+    needs_time = True
+    is_adversarial = False
+    folds_type_slice = False
+
+    def stage1_loss(self, model, x1, cond_cont, cond_cat, *, schedule=None) -> torch.Tensor:
+        return flow_matching_loss(model, x1, cond_cont, cond_cat)
+
+    def stage2_loss(
+        self,
+        model,
+        x1_s2,
+        cond_cont,
+        cond_cat,
+        stage1_ctx,
+        sec_mask,
+        *,
+        type_dim=None,
+        ar_inputs=None,
+    ) -> torch.Tensor:
+        if ar_inputs is not None:
+            return flow_matching_loss_secondary_ar(
+                model,
+                x1_s2,
+                cond_cont,
+                cond_cat,
+                stage1_ctx,
+                ar_inputs["history_feat"],
+                ar_inputs["has_prev"],
+                ar_inputs["remaining_frac"],
+                ar_inputs["slot_idx"],
+                sec_mask,
+                type_dim=type_dim,
+            )
+        return flow_matching_loss_secondary(model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=type_dim)
+
+
+@register_objective("ddpm")
+class DdpmObjective(Objective):
+    """Full DDPM ancestral sampling (Nichol & Dhariwal 2021 cosine schedule)
+    — the throwaway baseline. Stage-1 only: no `Stage2*` class has ever been
+    trained with `generator="ddpm"` in practice, so there's no stage-2 ddpm
+    loss to dispatch to (matches `FlowDDPMStageTrainer`'s pre-existing
+    stage-2 guard)."""
+
+    needs_time = True
+    is_adversarial = False
+    folds_type_slice = False
+    supports_stage2_decoder = False
+
+    def __init__(self, n_steps: int = 1000) -> None:
+        self.n_steps = n_steps
+
+    def build_schedule(self, n_steps: int, device: torch.device) -> CosineSchedule:
+        return CosineSchedule(T=n_steps).to(device)
+
+    def stage1_loss(self, model, x1, cond_cont, cond_cat, *, schedule=None) -> torch.Tensor:
+        assert schedule is not None, "DdpmObjective.stage1_loss needs a schedule (see build_schedule)"
+        return schedule.loss(model, x1, cond_cont, cond_cat)
+
+
+@register_objective("wgan")
+class WganObjective(Objective):
+    """WGAN-GP (Gulrajani et al. 2017) — single forward pass instead of an
+    ODE loop. `stage1_loss`/`stage2_loss` are unused: `WGANStageTrainer` owns
+    its own dual generator/critic step instead of a single scalar loss."""
+
+    needs_time = False
+    is_adversarial = True
+    folds_type_slice = True
+
+    def trunk_in_dim(self, out_dim: int, noise_dim: int) -> int:
+        return noise_dim

giant/model/routers.py

@@ -11,6 +11,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from giant.cond_layout import CondLayout
 from giant.constants import COND_DIM
 
 # ---------------------------------------------------------------------------
@@ -211,7 +212,7 @@ class PdgRouter(Router):
             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)
+        e = self.pdg_emb(cond_cat[:, CondLayout.PDG_COL])  # (B, emb_dim)
         d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(-1)  # (B, n_experts)
         return torch.softmax(-d2 / self.temperature, dim=-1)
 
@@ -243,8 +244,8 @@ class ProcessRouter(Router):
         )
 
     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])
+        pdg_e = self.pdg_emb(cond_cat[:, CondLayout.PDG_COL])
+        mat_e = self.mat_emb(cond_cat[:, CondLayout.MAT_COL])
         h = torch.cat([cond_cont, pdg_e, mat_e], dim=-1)
         return self.classifier(h)
 

giant/model/trunks.py

@@ -1,15 +1,61 @@
-"""Trunks: everything downstream of the fused conditioning vector — monolithic
-or expert-routed (issues.md Issue 8)."""
+"""Trunks: everything downstream of the fused conditioning vector — a
+registrable expert *body* architecture (`TRUNK_REGISTRY`/`register_trunk`),
+used standalone or mixed by a `Router` (issues.md Issue 8; trunk-selectability
+gitea #33).
+
+Whether a body is mixed is orthogonal to which body it is: `RoutedTrunk`
+builds `router.n_experts` instances of whichever body `trunk_type` names, so
+a future body (e.g. a transformer) automatically gets a mixture variant for
+free — no separate "routed transformer trunk" class needed.
+"""
 
 import torch
 import torch.nn as nn
 
-from giant.model.layers import ResBlock
+from giant.model.layers import build_block
 from giant.model.routers import Router
 
+TRUNK_REGISTRY: dict[str, type[nn.Module]] = {}
 
+
+def register_trunk(name: str):
+    def decorator(cls: type[nn.Module]) -> type[nn.Module]:
+        TRUNK_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def build_expert_body(
+    name: str,
+    in_dim: int,
+    out_dim: int,
+    hidden_dim: int,
+    n_blocks: int,
+    cond_dim: int,
+    dropout: float = 0.0,
+    block_conditioning: str = "add",
+) -> nn.Module:
+    """Factory: look up a registered trunk body by name and construct one
+    instance of it — used both for a standalone (unrouted) trunk and for each
+    expert inside a `RoutedTrunk`. `block_conditioning` selects the
+    `BLOCK_REGISTRY` entry each body's internal `ResBlock`-family blocks use
+    (`trunk.block_conditioning`, gitea #34) — an optional trailing kwarg a
+    future non-`ResBlock`-based body can simply ignore, same idiom as
+    `Trunk.forward`'s accept-and-ignore `cond_cont`/`cond_cat`."""
+    if name not in TRUNK_REGISTRY:
+        raise ValueError(f"unknown trunk type {name!r}; available: {sorted(TRUNK_REGISTRY)}")
+    cls = TRUNK_REGISTRY[name]
+    return cls(in_dim, out_dim, hidden_dim, n_blocks, cond_dim, dropout, block_conditioning=block_conditioning)
+
+
+@register_trunk("resmlp")
 class ExpertTrunk(nn.Module):
-    """One small expert: `input_proj -> ResBlock stack -> out_proj`.
+    """`input_proj -> ResBlock stack -> out_proj` — the registered `"resmlp"`
+    trunk body. Used both standalone (no router: `forward`'s `cond_cont`/
+    `cond_cat` are accepted and ignored, satisfying the `Trunk` interface
+    directly with no wrapper class) and as one expert inside a `RoutedTrunk`
+    (`_route_forward` calls it with just `(x, cond)`).
 
     Unlike v0.2, `out_dim` is independent of `in_dim` — needed by stage-2 AR
     tokens later (`noise_dim` in, `4 + type_dim` out), even though every
@@ -24,13 +70,23 @@ class ExpertTrunk(nn.Module):
         n_blocks: int,
         cond_dim: int,
         dropout: float = 0.0,
+        block_conditioning: str = "add",
     ) -> None:
         super().__init__()
+        self.out_dim = out_dim
         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.blocks = nn.ModuleList(
+            [build_block(block_conditioning, hidden_dim, cond_dim, dropout) for _ in range(n_blocks)]
+        )
         self.out_proj = nn.Linear(hidden_dim, out_dim)
 
-    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond: torch.Tensor,
+        cond_cont: torch.Tensor | None = None,
+        cond_cat: torch.Tensor | None = None,
+    ) -> torch.Tensor:
         x = self.input_proj(x)
         for block in self.blocks:
             x = block(x, cond)
@@ -55,13 +111,13 @@ def _route_forward(
     """
     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)
+        out = torch.zeros(x.shape[0], experts[0].out_dim, 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_dim = experts[0].out_dim
     out = torch.zeros(x.shape[0], out_dim, device=x.device)
     for i, expert in enumerate(experts):
         mask = idx == i
@@ -71,10 +127,10 @@ def _route_forward(
 
 
 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)."""
+    """Interface implemented by a standalone trunk body (any `TRUNK_REGISTRY`
+    entry, e.g. `ExpertTrunk`) and by `RoutedTrunk`: everything downstream of
+    the fused conditioning vector, i.e. the actual generative trunk of a
+    stage."""
 
     def forward(
         self,
@@ -86,49 +142,35 @@ class Trunk(nn.Module):
         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,
+        trunk_type: str,
         in_dim: int,
         out_dim: int,
         hidden_dim: int,
         n_res_blocks: int,
         cond_dim: int,
         dropout: float = 0.0,
+        block_conditioning: str = "add",
     ) -> None:
         super().__init__()
         self.router = router
         self.experts = nn.ModuleList(
-            [ExpertTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) for _ in range(router.n_experts)]
+            [
+                build_expert_body(
+                    trunk_type,
+                    in_dim,
+                    out_dim,
+                    hidden_dim,
+                    n_res_blocks,
+                    cond_dim,
+                    dropout,
+                    block_conditioning=block_conditioning,
+                )
+                for _ in range(router.n_experts)
+            ]
         )
 
     def forward(
@@ -143,13 +185,30 @@ class RoutedTrunk(Trunk):
 
 def build_trunk(
     router: Router | None,
+    trunk_type: str,
     in_dim: int,
     out_dim: int,
     hidden_dim: int,
     n_res_blocks: int,
     cond_dim: int,
     dropout: float = 0.0,
-) -> Trunk:
+    block_conditioning: str = "add",
+) -> nn.Module:
+    """Build a stage's trunk: `trunk_type` (a `TRUNK_REGISTRY` key, e.g.
+    `"resmlp"`) selects the expert body architecture; `router`, if given,
+    wraps `router.n_experts` instances of that body in a `RoutedTrunk`
+    mixture — otherwise a single body is returned directly (no wrapper
+    class), which is what makes an unrouted trunk's state-dict keys land
+    directly under `trunk.*` instead of `trunk.experts.0.*` (see
+    `giant.model._legacy.migrate_legacy_state_dict`, which assumes exactly
+    this flat layout for a v0.2 monolithic checkpoint). `block_conditioning`
+    (a `BLOCK_REGISTRY` key, e.g. `"add"`/`"film"`/`"adaln"`) selects each
+    body's conditioning-injection mechanism (gitea #34).
+    """
     if router is not None:
-        return RoutedTrunk(router, 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)
+        return RoutedTrunk(
+            router, trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout, block_conditioning
+        )
+    return build_expert_body(
+        trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout, block_conditioning
+    )

giant/pipeline.py

@@ -188,8 +188,8 @@ def run_setup_stage(
     # independent of both.
     particle_cfg = cfg["conditioning"]["particle"]
     material_cfg = cfg["conditioning"]["material"]
-    particle_type_cfg_dict = cfg["stage2_model"].get("particle_type") or {}
-    particle_type_target = config.ParticleTypeConfig.from_dict(particle_type_cfg_dict).target
+    particle_type_cfg = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type"))
+    particle_type_target = particle_type_cfg.target
 
     def _pdg_topn(n_classes: int) -> TopNMap:
         cache_key = setup_cache.topn_key("pdg", n_classes)
@@ -210,7 +210,7 @@ def run_setup_stage(
 
     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_n_classes = resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
         sec_type_topn_map = _pdg_topn(sec_type_n_classes)
 
     mat_topn_map: TopNMap | None = None
@@ -271,6 +271,13 @@ def run_setup_stage(
                     particle_conditioning=particle_conditioning,
                     material_conditioning=material_conditioning,
                     sec_phys_only=True,
+                    # This pass reads only cond_cont/sec_cont, never cond_cat —
+                    # but cond_cat's width is the conditioning modes' call
+                    # (giant.cond_layout.CondLayout), so an "onehot" axis still
+                    # has to be handed its map rather than silently yielding a
+                    # narrower array.
+                    pdg_topn_map=pdg_topn_map.class_map if pdg_topn_map is not None else None,
+                    mat_topn_map=mat_topn_map.class_map if mat_topn_map is not None else None,
                     k_max=k_max,
                 )
                 cond_cont = feats.cond_cont

giant/rollout.py

@@ -141,7 +141,7 @@ def decode_secondary_identity(
     Returns (sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg,
     sec_type_l1_dist) — the last is `None` except under `"embedding"`.
     """
-    target = sec_decoder.particle_type_cfg.get("target", "physical")
+    target = sec_decoder.particle_type_cfg.target
 
     if target == "physical":
         sec_full = torch.cat([sec_cont, sec_type], dim=-1).cpu().numpy()
@@ -491,7 +491,7 @@ 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:
+    if sec_decoder.particle_type_cfg.target == "onehot" and sec_type_topn_map is None:
         raise RuntimeError(
             "stage2_model.particle_type.target='onehot' rollout needs sec_type_topn_map "
             "(the checkpoint's saved top-N map) — see ckpt['sec_type_topn_map']"
@@ -681,7 +681,6 @@ def _step_chunk(
     post_pos = reconstruct_post_pos(tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local)
 
     n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cc, ck, stage1_norm, n_sec_pred_stage1)
-    n_sec_np = n_sec_pred.cpu().numpy().astype(np.int64)
 
     # --- Secondaries ---
     # No snapping for "physical"/history-facing state elsewhere in the
@@ -691,7 +690,13 @@ def _step_chunk(
     # decode_secondary_identity's docstring for how each
     # particle_type.target differs on whether PDG resolution is a real
     # identity decision or just a reporting label.
-    sec_cont, sec_type, _valid = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
+    sec_cont, sec_type, sec_valid = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
+    # A stop-token decoder resolves n_sec_pred=None above — the real count
+    # only exists once sample_stage2 has actually generated (or stopped
+    # generating) tokens, so read it back off sec_valid here. Under every
+    # other n_sec.mode sec_valid was built FROM n_sec_pred, so this is a
+    # no-op round trip in those cases.
+    n_sec_np = sec_valid.sum(dim=-1).cpu().numpy().astype(np.int64)
     sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, sec_type_l1_dist = decode_secondary_identity(
         sec_decoder,
         sec_cont,

giant/sample.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn.functional as F
 
 from giant.constants import CONT_SLOT_DIM, X_DIM
-from giant.model.network import Stage2Autoregressive, stage2_trunk_sec_dim
+from giant.model.network import DdpmObjective, Stage2Autoregressive, build_objective, stage2_trunk_sec_dim
 from giant.model.schedule import CosineSchedule
 
 
@@ -131,8 +131,8 @@ def _stage2_flat_width(sec_decoder: torch.nn.Module) -> int:
 
 
 def _type_folded(sec_decoder: torch.nn.Module) -> bool:
-    target = sec_decoder.particle_type_cfg.get("target", "physical")
-    return target == "physical" or sec_decoder.generator_kind == "wgan"
+    target = sec_decoder.particle_type_cfg.target
+    return target == "physical" or build_objective(sec_decoder.generator_kind).folds_type_slice
 
 
 def _decode_stage2_flat(
@@ -229,14 +229,14 @@ def sample_secondaries_ar(
     cond_cont: torch.Tensor,
     cond_cat: torch.Tensor,
     stage1_out: torch.Tensor,
-    n_sec_pred: torch.Tensor,
+    n_sec_pred: torch.Tensor | None,
     steps: int = 10,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """`Stage2Autoregressive` inference loop: one token at a time, in
-    descending-energy slot order, `k_max` sequential calls. Unlike training
-    (teacher forcing — a single parallel pass over ground-truth tokens, see
-    `giant.training.stage2_inputs._assemble_stage2_ar_inputs`), there is no
-    ground truth at inference: each token's conditioning is built
+    descending-energy slot order, up to `k_max` sequential calls. Unlike
+    training (teacher forcing — a single parallel pass over ground-truth
+    tokens, see `giant.training.stage2_inputs._assemble_stage2_ar_inputs`),
+    there is no ground truth at inference: each token's conditioning is built
     free-running, from the PREVIOUS TOKEN'S OWN just-generated output — the
     train/inference gap that is the cost of markov history's
     expressiveness.
@@ -244,7 +244,30 @@ def sample_secondaries_ar(
     A `{flow,ddpm}` token costs `steps` ODE substeps; `wgan` costs one pass —
     the "K sequential forwards" cost applies per-token here, not
     once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per
-    physics step.
+    physics step (or ~`n_sec * steps` under `n_sec_pred=None` below, once
+    every row in the batch has stopped).
+
+    `n_sec_pred`, if given, fixes each row's secondary count up front (as
+    resolved by `resolve_n_sec` — `n_sec.mode` in `("head", "truth")`, or a
+    stop-token decoder driven by `_assemble_stage2_ar_inputs_scheduled`'s
+    ground-truth `n_sec`, which must run the *full* `k_max`-length free-
+    running self-sample regardless of the decoder's own stop head — the
+    scheduled-sampling training contract does not truncate). This always
+    runs the full `k_max`-iteration loop, masking by the given count at the
+    end exactly as before.
+
+    `n_sec_pred=None` is only valid when `sec_decoder.stop_head` is set
+    (`n_sec.mode = "stop_token"`): before generating each slot's token, that
+    slot's own stop logit (`predict_stop`, evaluated on the same prefix
+    conditioning as the token itself — see `predict_type`'s docstring for
+    why this needs no extra state) decides whether generation should have
+    already stopped, per `sec_decoder.stop_sampling` ("greedy": threshold at
+    0; "sample": a Bernoulli draw at `sigmoid(logit)`). A row's own
+    `n_sec_pred` is the first slot index where this fires; once every row in
+    the batch has fired, the loop breaks before spending a model call on the
+    next slot's token — the average-case cost win the docstring above
+    describes. A row that never fires within `k_max` is capped there
+    (`K_MAX` stays a safety cap, not a modeling ceiling).
 
     Under `history="attention"` the history encoding is computed once per
     slot via `Stage2Autoregressive.history_step` (a KV-cache append)
@@ -281,8 +304,8 @@ def sample_secondaries_ar(
     device = cond_cont.device
     k_max = sec_decoder.k_max
     type_dim = sec_decoder.type_dim
-    generator = sec_decoder.generator_kind
-    target = sec_decoder.particle_type_cfg.get("target", "physical")
+    objective = build_objective(sec_decoder.generator_kind)
+    target = sec_decoder.particle_type_cfg.target
     type_folded = _type_folded(sec_decoder)
     token_dim = CONT_SLOT_DIM + type_dim if type_folded else CONT_SLOT_DIM
 
@@ -294,6 +317,16 @@ def sample_secondaries_ar(
     remaining = torch.ones(B, device=device)
     history_cache = sec_decoder.init_history_cache()
 
+    use_stop_token = n_sec_pred is None
+    if use_stop_token:
+        assert getattr(sec_decoder, "stop_head", None) is not None, (
+            "sample_secondaries_ar called with n_sec_pred=None on a decoder "
+            "with no stop_head — only valid under stage2_model.n_sec.mode = "
+            "'stop_token'"
+        )
+        finished = torch.zeros(B, dtype=torch.bool, device=device)
+        derived_n_sec = torch.full((B,), k_max, dtype=torch.long, device=device)
+
     for k in range(k_max):
         has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device)
         history_feat = prev_repr.unsqueeze(1)  # (B, 1, CONT_SLOT_DIM + type_dim)
@@ -301,7 +334,27 @@ def sample_secondaries_ar(
         slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
         hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache)
 
-        if generator == "wgan":
+        if use_stop_token:
+            stop_logit = sec_decoder.predict_stop(
+                cond_cont,
+                cond_cat,
+                stage1_out,
+                history_feat,
+                has_prev,
+                remaining_frac,
+                slot_idx,
+                hist=hist,
+            ).squeeze(1)
+            if sec_decoder.stop_sampling == "sample":
+                stop_now = torch.rand(B, device=device) < torch.sigmoid(stop_logit)
+            else:
+                stop_now = stop_logit >= 0.0
+            derived_n_sec[stop_now & ~finished] = k
+            finished = finished | stop_now
+            if finished.all():
+                break
+
+        if objective.is_adversarial:
             z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
             token = sec_decoder(
                 z,
@@ -362,7 +415,8 @@ def sample_secondaries_ar(
         prev_repr = torch.cat([stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1)
         remaining = torch.clamp(remaining * (1.0 - stick_fraction), min=0.0)
 
-    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1)
+    resolved_n_sec = derived_n_sec if use_stop_token else n_sec_pred
+    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < resolved_n_sec.unsqueeze(1)
     return sec_cont, sec_type, sec_valid
 
 
@@ -383,10 +437,10 @@ def sample_stage1(
     ddpm_steps: int = 1000,
 ) -> tuple[torch.Tensor, torch.Tensor | None]:
     """Dispatches on `stage1_model.generator_kind`."""
-    kind = stage1_model.generator_kind
-    if kind == "wgan":
+    objective = build_objective(stage1_model.generator_kind)
+    if objective.is_adversarial:
         return sample_wgan(stage1_model, cond_cont, cond_cat)
-    if kind == "ddpm":
+    if isinstance(objective, DdpmObjective):
         schedule = CosineSchedule(T=ddpm_steps).to(cond_cont.device)
         return sample_ddpm(stage1_model, cond_cont, cond_cat, schedule)
     return sample_flow(stage1_model, cond_cont, cond_cat, steps=steps)
@@ -397,7 +451,7 @@ def sample_stage2(
     cond_cont: torch.Tensor,
     cond_cat: torch.Tensor,
     stage1_out: torch.Tensor,
-    n_sec_pred: torch.Tensor,
+    n_sec_pred: torch.Tensor | None,
     steps: int,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Dispatches on `decoder` (one-shot vs autoregressive — the class
@@ -406,10 +460,16 @@ def sample_stage2(
     built with `generator="ddpm"` in practice and `flow_matching_loss_secondary*`
     is the only stage-2 training path that exists for the non-adversarial
     case, so there's nothing to dispatch to here.
+
+    `n_sec_pred=None` (from `resolve_n_sec` on a stop-token decoder) is only
+    meaningful for the autoregressive path — see `sample_secondaries_ar`'s
+    docstring; the one-shot samplers have no per-token stop mechanism to
+    derive a count from, so `n_sec_pred` must already be resolved for them.
     """
     if isinstance(sec_decoder, Stage2Autoregressive):
         return sample_secondaries_ar(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
-    if sec_decoder.generator_kind == "wgan":
+    assert n_sec_pred is not None, "one-shot stage-2 decoders need a resolved n_sec_pred"
+    if build_objective(sec_decoder.generator_kind).is_adversarial:
         return sample_secondaries_wgan(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred)
     return sample_secondaries(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
 
@@ -421,20 +481,31 @@ def resolve_n_sec(
     cond_cat: torch.Tensor,
     stage1_out: torch.Tensor,
     n_sec_pred: torch.Tensor | None,
-) -> torch.Tensor:
+) -> torch.Tensor | None:
     """`n_sec_pred` is already populated when `stage1_model` owns a legacy
     `n_sec_head` (a migrated v0.2 checkpoint — see `Stage1Model`'s
-    docstring); otherwise ask stage 2, which owns it by default. Raises if
-    neither stage owns a head at all — the only way that happens is
-    `stage2_model.n_sec.mode` other than `"head"` (`"truth"`/`"stop_token"`),
-    neither of which is a valid rollout-/predict-capable checkpoint."""
+    docstring); otherwise ask stage 2, which owns it by default.
+
+    Returns `None` when `sec_decoder` owns a `stop_head` (`n_sec.mode =
+    "stop_token"`) instead of an `n_sec_head` — there is nothing to resolve
+    up front in that case, since the count only exists once
+    `sample_secondaries_ar` has actually generated (or stopped generating)
+    tokens; the caller passes this `None` straight through to `sample_stage2`
+    and reads the real count back off its returned `sec_valid`
+    (`sec_valid.sum(-1)`) afterwards.
+
+    Raises if neither stage owns any n_sec mechanism at all — the only way
+    that happens is `stage2_model.n_sec.mode = "truth"`, which is not a valid
+    rollout-/predict-capable checkpoint."""
     if n_sec_pred is not None:
         return n_sec_pred
+    if getattr(sec_decoder, "stop_head", None) is not None:
+        return None
     if getattr(sec_decoder, "n_sec_head", None) is None:
         raise RuntimeError(
-            "checkpoint has no n_sec_head on either stage — needs "
-            "stage2_model.n_sec.mode = 'head' (the default); 'truth' is "
-            "standalone-evaluation-only and 'stop_token' isn't implemented"
+            "checkpoint has no n_sec_head/stop_head on either stage — needs "
+            "stage2_model.n_sec.mode = 'head' (the default) or 'stop_token'; "
+            "'truth' is standalone-evaluation-only"
         )
     logits = sec_decoder.predict_n_sec(cond_cont, cond_cat, stage1_out)
     return logits.argmax(dim=-1)

giant/tools/dwarf.py

@@ -442,43 +442,65 @@ def warm_cache(
         Path,
         typer.Argument(help="Parquet file, directory, or .manifest — same as `giant train`'s"),
     ],
+    config: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--config",
+            "-c",
+            help="TOML config file to warm for — same file the `giant train` run(s) will use. "
+            "Mutually exclusive with the flags below (put val-fraction/seed/conditioning/router "
+            "settings in the file itself, so warming and training can't disagree on them)",
+        ),
+    ] = None,
     val_fraction: Annotated[
-        float,
+        Optional[float],
         typer.Option(
             "--val-fraction",
             "-f",
-            help="Must match the `giant train` run(s) to warm for",
+            help="Must match the `giant train` run(s) to warm for. Not allowed together with --config",
         ),
-    ] = 0.1,
+    ] = None,
     seed: Annotated[
-        int,
-        typer.Option("--seed", "-s", help="Must match the `giant train` run(s) to warm for"),
-    ] = 0,
+        Optional[int],
+        typer.Option(
+            "--seed",
+            "-s",
+            help="Must match the `giant train` run(s) to warm for. Not allowed together with --config",
+        ),
+    ] = None,
     particle_conditioning: Annotated[
-        Conditioning,
+        Optional[Conditioning],
         typer.Option(
             "--particle-conditioning",
-            help="Must match the `giant train` run(s)' conditioning.particle.type to warm for",
+            help="Must match the `giant train` run(s)' conditioning.particle.type to warm for. "
+            "Not allowed together with --config",
         ),
-    ] = Conditioning.physical,
+    ] = None,
     material_conditioning: Annotated[
-        Conditioning,
+        Optional[Conditioning],
         typer.Option(
             "--material-conditioning",
             help="Must match the `giant train` run(s)' conditioning.material.type "
             "to warm for — independent of --particle-conditioning "
-            "(the two axes may differ)",
+            "(the two axes may differ). Not allowed together with --config",
         ),
-    ] = Conditioning.physical,
+    ] = None,
     router: Annotated[
-        bool,
+        Optional[bool],
         typer.Option(
             "--router/--no-router",
-            help="Warm the process vocabulary too (only takes effect with --router-type process)",
+            help="Warm the process vocabulary too (only takes effect with --router-type process). "
+            "Not allowed together with --config",
         ),
-    ] = False,
-    router_type: Annotated[str, typer.Option("--router-type", help="Router implementation name")] = "energy",
-    n_experts: Annotated[int, typer.Option("--n-experts", help="Number of routed experts")] = 4,
+    ] = None,
+    router_type: Annotated[
+        Optional[str],
+        typer.Option("--router-type", help="Router implementation name. Not allowed together with --config"),
+    ] = None,
+    n_experts: Annotated[
+        Optional[int],
+        typer.Option("--n-experts", help="Number of routed experts. Not allowed together with --config"),
+    ] = None,
     rebuild: Annotated[
         bool,
         typer.Option("--rebuild", help="Ignore any existing sidecar and recompute every section"),
@@ -487,17 +509,38 @@ def warm_cache(
     """Precompute `giant train`'s setup-stage sidecar for `data` ahead of time.
 
     Warms the vocab maps, event-id split index, and the normalizer entry for
-    the given --val-fraction/--seed/--particle-conditioning/
-    --material-conditioning, so a later `giant train` run (or a `dwarf
-    hparam-scan` sweep, which shares one such entry across every run) skips
-    straight to training. See giant/data/setup_cache.py.
+    either --config, or the given --val-fraction/--seed/
+    --particle-conditioning/--material-conditioning/--router* flags, so a
+    later `giant train` run (or a `dwarf hparam-scan` sweep, which shares one
+    such entry across every run) skips straight to training. See
+    giant/data/setup_cache.py.
     """
+    flag_overrides = {
+        "--val-fraction": val_fraction,
+        "--seed": seed,
+        "--particle-conditioning": particle_conditioning,
+        "--material-conditioning": material_conditioning,
+        "--router/--no-router": router,
+        "--router-type": router_type,
+        "--n-experts": n_experts,
+    }
+    if config is not None:
+        given = [name for name, value in flag_overrides.items() if value is not None]
+        if given:
+            typer.echo(
+                f"error: --config cannot be combined with {', '.join(given)} "
+                "— put these settings in the config file instead",
+                err=True,
+            )
+            raise typer.Exit(1)
+
     run_warm_setup_cache(
         data=str(data),
+        config_path=config,
         val_fraction=val_fraction,
         seed=seed,
-        particle_conditioning=particle_conditioning.value,
-        material_conditioning=material_conditioning.value,
+        particle_conditioning=particle_conditioning.value if particle_conditioning is not None else None,
+        material_conditioning=material_conditioning.value if material_conditioning is not None else None,
         router_enabled=router,
         router_type=router_type,
         n_experts=n_experts,

giant/tools/warm_setup_cache.py

@@ -10,63 +10,86 @@ for the sidecar itself.
 from pathlib import Path
 
 from giant import config as gconfig
-from giant.constants import K_MAX
 from giant.pipeline import run_setup_stage
 
 
 def run_warm_setup_cache(
     data: str,
-    val_fraction: float = 0.1,
-    seed: int = 0,
-    particle_conditioning: str = "physical",
-    material_conditioning: str = "physical",
-    router_enabled: bool = False,
-    router_type: str = "energy",
-    n_experts: int = 4,
+    config_path: Path | None = None,
+    val_fraction: float | None = None,
+    seed: int | None = None,
+    particle_conditioning: str | None = None,
+    material_conditioning: str | None = None,
+    router_enabled: bool | None = None,
+    router_type: str | None = None,
+    n_experts: int | None = None,
     rebuild: bool = False,
     echo=print,
 ) -> None:
     """Populate (or refresh) the setup cache sidecar for `data`.
 
-    `val_fraction`/`seed`/`particle_conditioning`/`material_conditioning`
-    select the normalizer cache entry
-    (`giant.data.setup_cache.normalizer_key`) — pass the same values a later
-    `giant train` invocation will use so it hits this warmed entry. The two
-    conditioning axes are independent and may differ.
-    `router_enabled`/`router_type`/`n_experts` only matter for
-    `router_type == "process"` (warms that `n_experts`'s process map); the
-    energy-router quantile summary is always collected regardless, so a
-    later `--router-type energy` run never needs to rescan just to seed
-    centers.
+    Two mutually exclusive ways to select what to warm for (enforced by the
+    caller, `giant.tools.dwarf.warm_cache` — this function just trusts
+    whichever combination it's given):
+
+    - `config_path`: the same TOML `giant train --config` takes. Every value
+      `run_setup_stage` needs (`train.val_fraction`/`seed`,
+      `conditioning.particle`/`material.type`, both stages' `router`,
+      `stage2_model.particle_type.n_classes`, ...) is read from the one
+      resulting merged `cfg`, so a later `giant train --config <same file>`
+      run resolves to exactly the same cache keys — see gitea #59.
+    - The individual flags below: `val_fraction`/`seed`/
+      `particle_conditioning`/`material_conditioning` select the normalizer
+      cache entry (`giant.data.setup_cache.normalizer_key`) — pass the same
+      values a later `giant train` invocation will use so it hits this
+      warmed entry. The two conditioning axes are independent and may
+      differ. `router_enabled`/`router_type`/`n_experts` only matter for
+      `router_type == "process"` (warms that `n_experts`'s process map); the
+      energy-router quantile summary is always collected regardless, so a
+      later `--router-type energy` run never needs to rescan just to seed
+      centers.
+
+    Any flag left `None` is omitted from the merge, so it falls back to
+    `DEFAULT_CONFIG`'s own value (or the config file's, if `config_path` is
+    given) instead of silently overriding it — see gitea #59.
     """
-    router_cfg = {
-        "enabled": router_enabled,
-        "type": router_type,
-        "n_experts": n_experts,
-    }
-    # Merged against DEFAULT_CONFIG (not a hand-rolled partial dict) so
-    # run_setup_stage always sees every key it might read (e.g.
-    # conditioning.particle.emb_dim, stage2_model.particle_type.target) at
-    # its real default, not silently missing/None — see issues.md Issue 1.
+    overrides: dict = {}
+
+    conditioning_overrides: dict = {}
+    if particle_conditioning is not None:
+        conditioning_overrides["particle"] = {"type": particle_conditioning}
+    if material_conditioning is not None:
+        conditioning_overrides["material"] = {"type": material_conditioning}
+    if conditioning_overrides:
+        overrides["conditioning"] = conditioning_overrides
+
     # This CLI only ever configures one router (matching today's single
-    # --router-type flag), so it's placed on stage1_model; stage2_model's
-    # stays disabled.
-    cfg = gconfig.merge_cli_overrides(
-        gconfig.DEFAULT_CONFIG,
-        None,
-        {
-            "conditioning": {
-                "particle": {"type": particle_conditioning},
-                "material": {"type": material_conditioning},
-            },
-            "stage1_model": {"router": router_cfg},
-            "stage2_model": {"router": {"enabled": False}, "k_max": K_MAX},
-        },
-    )
+    # --router-type flag), so it's placed on stage1_model; stage2_model's is
+    # left to DEFAULT_CONFIG/the config file rather than forced disabled.
+    router_overrides: dict = {}
+    if router_enabled is not None:
+        router_overrides["enabled"] = router_enabled
+    if router_type is not None:
+        router_overrides["type"] = router_type
+    if n_experts is not None:
+        router_overrides["n_experts"] = n_experts
+    if router_overrides:
+        overrides["stage1_model"] = {"router": router_overrides}
+
+    train_overrides: dict = {}
+    if val_fraction is not None:
+        train_overrides["val_fraction"] = val_fraction
+    if seed is not None:
+        train_overrides["seed"] = seed
+    if train_overrides:
+        overrides["train"] = train_overrides
+
+    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config_path, overrides)
+    gconfig.validate_config(cfg)
     run_setup_stage(
         Path(data),
-        val_fraction=val_fraction,
-        seed=seed,
+        val_fraction=cfg["train"]["val_fraction"],
+        seed=cfg["train"]["seed"],
         cfg=cfg,
         cache_setup=True,
         rebuild_setup_cache=rebuild,

giant/training/stage2_inputs.py

@@ -11,7 +11,9 @@ live in one place and stay unit-testable on their own.
 import torch
 import torch.nn.functional as F
 
+from giant.config import ParticleTypeConfig
 from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM
+from giant.model.objectives import build_objective
 from giant.sample import sample_secondaries_ar
 
 
@@ -30,7 +32,7 @@ def _gumbel_tau(step: int, total_steps: int, tau_start: float, tau_end: float) -
 def _type_repr(
     sec_type_idx: torch.Tensor,
     sec_cont: torch.Tensor,
-    particle_type_cfg: dict,
+    particle_type_cfg: ParticleTypeConfig,
     cond_enc: torch.nn.Module,
     emb_dim: int,
 ) -> torch.Tensor:
@@ -45,7 +47,7 @@ def _type_repr(
     latter must always reflect the true physical secondary that came before,
     regardless of what the *current* token's own training objective is.
     """
-    target = particle_type_cfg.get("target", "physical")
+    target = particle_type_cfg.target
     if target == "physical":
         return sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM]
     if target == "onehot":
@@ -56,7 +58,7 @@ def _type_repr(
 def _assemble_stage2_ar_target(
     sec_cont: torch.Tensor,
     sec_type_idx: torch.Tensor,
-    particle_type_cfg: dict,
+    particle_type_cfg: ParticleTypeConfig,
     generator: str,
     cond_enc: torch.nn.Module,
     emb_dim: int,
@@ -69,19 +71,20 @@ def _assemble_stage2_ar_target(
 
     - `target = "physical"`: unchanged from v0.2 — `sec_cont` (stick_logit,
       dir, log_mass, charge) as-is.
-    - `target` in `("onehot", "embedding")` + `generator in ("flow", "ddpm")`:
-      just the continuous stick/dir slots — the type slice isn't part of
-      this tensor at all (`type_head` handles it separately).
-    - `target` in `("onehot", "embedding")` + `generator == "wgan"`: stick/dir
-      slots concatenated with the per-slot type representation (a one-hot of
-      the true class, relaxed on the *generated* side only, by the caller;
-      or the conditioning's own detached embedding-table row).
+    - `target` in `("onehot", "embedding")` + an objective that doesn't fold
+      the type slice (flow/ddpm): just the continuous stick/dir slots — the
+      type slice isn't part of this tensor at all (`type_head` handles it
+      separately).
+    - `target` in `("onehot", "embedding")` + a folding objective (wgan):
+      stick/dir slots concatenated with the per-slot type representation (a
+      one-hot of the true class, relaxed on the *generated* side only, by the
+      caller; or the conditioning's own detached embedding-table row).
     """
-    target = particle_type_cfg.get("target", "physical")
+    target = particle_type_cfg.target
     if target == "physical":
         return sec_cont
     cont = sec_cont[..., :CONT_SLOT_DIM]
-    if generator != "wgan":
+    if not build_objective(generator).folds_type_slice:
         return cont
     type_repr = _type_repr(sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim)
     return torch.cat([cont, type_repr], dim=-1)
@@ -90,7 +93,7 @@ def _assemble_stage2_ar_target(
 def _assemble_stage2_real(
     sec_cont: torch.Tensor,
     sec_type_idx: torch.Tensor,
-    particle_type_cfg: dict,
+    particle_type_cfg: ParticleTypeConfig,
     generator: str,
     cond_enc: torch.nn.Module,
     emb_dim: int,
@@ -137,6 +140,27 @@ def _ar_has_prev(k_max: int, device: torch.device) -> torch.Tensor:
     return (torch.arange(k_max, device=device) >= 1).unsqueeze(0)
 
 
+def _stop_target_and_mask(n_sec: torch.Tensor, k_max: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+    """`(target, mask)`, both `(B, K_MAX)`, for `n_sec.mode = "stop_token"`'s
+    per-slot EOS head (`Stage2Autoregressive.predict_stop`).
+
+    `predict_stop` is evaluated on slot `k`'s own (pre-token) conditioning —
+    "should generation have already stopped by here" — so `target[k] = 1`
+    exactly at `k == n_sec` (the first invalid slot: `sample_secondaries_ar`
+    checks this before spending a model call generating that slot's token),
+    `0` elsewhere. `mask` is `k <= n_sec` — one slot *wider* than
+    `StageTrainer._sec_mask`'s `k < n_sec` token-content mask, since the stop
+    slot itself (`k == n_sec`) must be supervised even though there is no
+    real secondary there. A row with `n_sec == k_max` has no in-range stop
+    slot at all: `mask` covers the full `k_max` range (every generated token
+    is real) and `target` is all-zero — `sample_secondaries_ar` correctly
+    never breaks early for it, running into the `k_max` safety cap instead."""
+    idx = torch.arange(k_max, device=device).unsqueeze(0)
+    target = (idx == n_sec.unsqueeze(1)).float()
+    mask = idx <= n_sec.unsqueeze(1)
+    return target, mask
+
+
 def _ar_meta(k_max: int, batch: int, device: torch.device, fraction: torch.Tensor) -> dict[str, torch.Tensor]:
     """`has_prev`/`remaining_frac`/`slot_idx` — the three per-token AR
     conditioning tensors that don't depend on *which* history representation
@@ -155,7 +179,7 @@ def _ar_meta(k_max: int, batch: int, device: torch.device, fraction: torch.Tenso
 def _assemble_stage2_ar_inputs(
     sec_cont: torch.Tensor,
     sec_type_idx: torch.Tensor,
-    particle_type_cfg: dict,
+    particle_type_cfg: ParticleTypeConfig,
     cond_enc: torch.nn.Module,
     emb_dim: int,
 ) -> dict[str, torch.Tensor]:
@@ -198,7 +222,7 @@ def _stage2_tf_prob(mode: str, p_start: float, p_end: float, epoch: int, total_e
 def _history_repr_from_ar_sample(
     sec_cont_pred: torch.Tensor,
     sec_type_pred: torch.Tensor,
-    particle_type_cfg: dict,
+    particle_type_cfg: ParticleTypeConfig,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """`(fraction, direction, type_repr)` — the same triple `_type_repr` /
     `_stick_fraction` derive from ground truth, but from a free-running
@@ -211,7 +235,7 @@ def _history_repr_from_ar_sample(
     representation."""
     fraction = torch.sigmoid(sec_cont_pred[..., 0])
     direction = sec_cont_pred[..., 1:4]
-    if particle_type_cfg.get("target", "physical") == "onehot":
+    if particle_type_cfg.target == "onehot":
         type_dim = sec_type_pred.size(-1)
         type_repr = F.one_hot(sec_type_pred.argmax(-1), num_classes=type_dim).float()
     else:
@@ -227,7 +251,7 @@ def _assemble_stage2_ar_inputs_scheduled(
     sec_cont: torch.Tensor,
     sec_type_idx: torch.Tensor,
     n_sec: torch.Tensor,
-    particle_type_cfg: dict,
+    particle_type_cfg: ParticleTypeConfig,
     cond_enc: torch.nn.Module,
     emb_dim: int,
     p_tf: float,

giant/training/trainers.py

@@ -25,13 +25,7 @@ import torch.optim as optim
 from giant.config import ParticleTypeConfig, Stage1ModelConfig, Stage2ModelConfig, TrainConfig
 from giant.constants import CONT_SLOT_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,
-    flow_matching_loss_secondary,
-    flow_matching_loss_secondary_ar,
-)
+from giant.model.network import Router, build_objective, resolve_type_n_classes, stage2_type_dim
 from giant.model.wgan import generator_loss, gradient_penalty
 from giant.training.metrics import MetricSpec, stage_metric, train_metric, val_metric
 from giant.training.stage2_inputs import (
@@ -40,6 +34,7 @@ from giant.training.stage2_inputs import (
     _gumbel_tau,
     _relax_onehot_type_slice,
     _stage2_tf_prob,
+    _stop_target_and_mask,
 )
 
 
@@ -95,6 +90,7 @@ class StageSpec:
     # loss weights
     lambda_weight: float = 1.0
     n_sec_lambda: float = 0.1
+    n_sec_mode: str = "head"
 
     # particle-type target (stage 2 only)
     particle_type: ParticleTypeConfig = field(default_factory=ParticleTypeConfig)
@@ -150,9 +146,10 @@ class StageSpec:
             decoder=s2_spec.decoder if is_stage2 else "one_shot",
             lambda_weight=stage_spec.lambda_weight,
             n_sec_lambda=s2_spec.n_sec.lambda_weight,
+            n_sec_mode=s2_spec.n_sec.mode,
             particle_type=s2_spec.particle_type,
             particle_type_n_classes=resolve_type_n_classes(
-                s2_spec.particle_type.to_dict(), cfg["conditioning"]["particle"]["emb_dim"]
+                s2_spec.particle_type, 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
@@ -239,7 +236,7 @@ class StageTrainer:
         self.router = _stage_router(self.model)
         self._modules = (self.model, *extra_modules)
 
-        self.particle_type_cfg = spec.particle_type.to_dict()
+        self.particle_type_cfg = spec.particle_type
         self.particle_type_n_classes = spec.particle_type_n_classes
         self.ema_decay = spec.ema_decay
 
@@ -380,10 +377,9 @@ class StageTrainer:
         stage1-vs-stage2 `predict_n_sec` signature split, shared by the
         non-adversarial and WGAN trainers.
 
-        Gated on `n_sec_head is None`, not on `n_sec.mode`: a future
-        `mode="stop_token"` model (currently rejected in
-        `validate_config`) carries no head and would train its EOS signal in
-        the generator/AR loss path instead, so this correctly stays zero.
+        Gated on `n_sec_head is None`, not on `n_sec.mode`: a `mode =
+        "stop_token"` model carries no head at all (see `_stop_loss` for its
+        EOS signal instead), so this correctly stays zero for it.
         """
         if self.model.n_sec_head is None:
             zero = torch.zeros((), device=device)
@@ -397,6 +393,47 @@ class StageTrainer:
         nsec_acc = (logits.argmax(dim=-1) == n_sec).float().mean()
         return l_nsec, nsec_acc
 
+    def _stop_loss(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_ctx: torch.Tensor,
+        n_sec: torch.Tensor,
+        device: torch.device,
+        ar_inputs: dict[str, torch.Tensor] | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """`(l_stop, stop_acc)` for `n_sec.mode = "stop_token"`'s per-slot EOS
+        head (`Stage2Autoregressive.predict_stop`) — zeros when this stage
+        owns no `stop_head` (every other `n_sec.mode`), the same gating
+        convention `_n_sec_loss` uses for `n_sec_head`. The two heads are
+        mutually exclusive (`giant.model.builders`), so exactly one of
+        `_n_sec_loss`/`_stop_loss` is ever non-zero for a given stage.
+
+        Masked BCE against `_stop_target_and_mask`'s per-slot target — one
+        slot wider than `sec_mask` (the stop slot itself, `k == n_sec`, needs
+        supervision even though it holds no real secondary)."""
+        stop_head = getattr(self.model, "stop_head", None)
+        if stop_head is None:
+            zero = torch.zeros((), device=device)
+            return zero, zero
+        assert ar_inputs is not None
+        logits = self.model.predict_stop(
+            cond_cont,
+            cond_cat,
+            stage1_ctx,
+            ar_inputs["history_feat"],
+            ar_inputs["has_prev"],
+            ar_inputs["remaining_frac"],
+            ar_inputs["slot_idx"],
+        )
+        target, mask = _stop_target_and_mask(n_sec, logits.size(1), device)
+        mask_f = mask.float()
+        denom = mask_f.sum().clamp(min=1)
+        bce = F.binary_cross_entropy_with_logits(logits, target, reduction="none")
+        l_stop = (bce * mask_f).sum() / denom
+        stop_acc = (((logits >= 0).float() == target).float() * mask_f).sum() / denom
+        return l_stop, stop_acc
+
     @staticmethod
     def _step_optimizer(optimizer: optim.Optimizer, loss: torch.Tensor, params: list) -> float:
         """`zero_grad -> backward -> clip_grad_norm_(1.0) -> step`, returning
@@ -450,20 +487,23 @@ class FlowDDPMStageTrainer(StageTrainer):
     """flow or ddpm generator for a single stage."""
 
     def __init__(self, spec: StageSpec, model: torch.nn.Module, device: torch.device) -> None:
-        if spec.is_stage2 and spec.generator not in ("flow",):
+        objective = build_objective(spec.generator, n_steps=spec.ddpm_n_steps)
+        if spec.is_stage2 and not objective.supports_stage2_decoder:
             raise NotImplementedError(
                 f"stage2_model.generator={spec.generator!r} is accepted by the "
                 "schema but not implemented in v0.3.0 for stage 2 (only "
                 "'flow' and 'wgan' have a stage-2 secondary-decoder loss)"
             )
         super().__init__(spec, model, device)
-        self.particle_type_lambda = self.particle_type_cfg.get("lambda", 1.0)
+        self.objective = objective
+        self.particle_type_lambda = self.particle_type_cfg.lambda_weight
         # Width of the type slice actually folded into x1_s2 by _sec_target,
-        # under this trainer's generator (flow/ddpm only — see the
-        # NotImplementedError above): "physical" keeps it folded in
-        # (PARTICLE_PHYS_DIM wide, unchanged from v0.2); "onehot"/"embedding"
-        # pull it out into model.type_head instead (0 here).
-        self._flow_type_dim = None if self.particle_type_cfg.get("target", "physical") == "physical" else 0
+        # under this trainer's objective (flow/ddpm only — see the
+        # NotImplementedError above, neither folds the type slice): "physical"
+        # keeps it folded in (PARTICLE_PHYS_DIM wide, unchanged from v0.2);
+        # "onehot"/"embedding" pull it out into model.type_head instead (0
+        # here).
+        self._flow_type_dim = None if self.particle_type_cfg.target == "physical" else 0
 
         self.params = list(self.model.parameters())
         self.optimizer = optim.AdamW(self.params, lr=spec.lr, weight_decay=spec.weight_decay)
@@ -472,7 +512,7 @@ class FlowDDPMStageTrainer(StageTrainer):
             warmup_steps=spec.warmup_epochs * spec.steps_per_epoch,
             total_steps=max(spec.epochs * spec.steps_per_epoch, 1),
         )
-        self.ddpm_schedule = CosineSchedule(T=spec.ddpm_n_steps).to(device) if spec.generator == "ddpm" else None
+        self.ddpm_schedule = self.objective.build_schedule(spec.ddpm_n_steps, device)
 
         self.train_metrics = [
             train_metric(key)
@@ -480,10 +520,12 @@ class FlowDDPMStageTrainer(StageTrainer):
                 "loss",
                 "loss_gen",
                 "loss_nsec",
+                "loss_stop",
                 "loss_balance",
                 "loss_proc",
                 "loss_entropy",
                 "nsec_acc",
+                "stop_acc",
                 "loss_type",
                 "type_acc",
                 "grad_norm",
@@ -496,6 +538,8 @@ class FlowDDPMStageTrainer(StageTrainer):
                 "loss_gen",
                 "loss_nsec",
                 "nsec_acc",
+                "loss_stop",
+                "stop_acc",
                 "loss_type",
                 "type_acc",
             )
@@ -504,26 +548,9 @@ class FlowDDPMStageTrainer(StageTrainer):
 
     def _generator_loss(self, cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=None):
         if not self.is_stage2:
-            if self.generator == "flow":
-                return flow_matching_loss(self.model, x1_s1, cond_cont, cond_cat)
-            assert self.ddpm_schedule is not None
-            return self.ddpm_schedule.loss(self.model, x1_s1, cond_cont, cond_cat)
-        if self.decoder == "autoregressive":
-            assert ar_inputs is not None
-            return flow_matching_loss_secondary_ar(
-                self.model,
-                x1_s2,
-                cond_cont,
-                cond_cat,
-                stage1_ctx,
-                ar_inputs["history_feat"],
-                ar_inputs["has_prev"],
-                ar_inputs["remaining_frac"],
-                ar_inputs["slot_idx"],
-                sec_mask,
-                type_dim=self._flow_type_dim,
-            )
-        return flow_matching_loss_secondary(
+            return self.objective.stage1_loss(self.model, x1_s1, cond_cont, cond_cat, schedule=self.ddpm_schedule)
+        assert self.decoder != "autoregressive" or ar_inputs is not None
+        return self.objective.stage2_loss(
             self.model,
             x1_s2,
             cond_cont,
@@ -531,6 +558,7 @@ class FlowDDPMStageTrainer(StageTrainer):
             stage1_ctx,
             sec_mask,
             type_dim=self._flow_type_dim,
+            ar_inputs=ar_inputs,
         )
 
     def _type_loss(
@@ -568,7 +596,7 @@ class FlowDDPMStageTrainer(StageTrainer):
             type_out = self.model.predict_type(cond_cont, cond_cat, stage1_ctx)
         mask = sec_mask.float()
         denom = mask.sum().clamp(min=1)
-        if self.particle_type_cfg.get("target") == "onehot":
+        if self.particle_type_cfg.target == "onehot":
             ce = F.cross_entropy(type_out.transpose(1, 2), sec_type_idx, reduction="none")
             l_type = (ce * mask).sum() / denom
             type_acc = ((type_out.argmax(-1) == sec_type_idx).float() * mask).sum() / denom
@@ -605,6 +633,7 @@ class FlowDDPMStageTrainer(StageTrainer):
 
         l_gen = self._generator_loss(cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=ar_inputs)
         l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device)
+        l_stop, stop_acc = self._stop_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device, ar_inputs)
 
         l_type, type_acc = self._type_loss(
             cond_cont,
@@ -625,7 +654,11 @@ class FlowDDPMStageTrainer(StageTrainer):
             if self.spec.lambda_entropy > 0:
                 l_entropy = self.router.entropy_loss(cond_cont, cond_cat)
 
-        total = self.spec.lambda_weight * l_gen + self.spec.n_sec_lambda * l_nsec + self.particle_type_lambda * l_type
+        total = (
+            self.spec.lambda_weight * l_gen
+            + self.spec.n_sec_lambda * (l_nsec + l_stop)
+            + self.particle_type_lambda * l_type
+        )
         if self.spec.lambda_balance > 0:
             total = total + self.spec.lambda_balance * l_balance
         if self.spec.lambda_proc > 0:
@@ -637,12 +670,14 @@ class FlowDDPMStageTrainer(StageTrainer):
             "loss": total,
             "loss_gen": l_gen,
             "loss_nsec": l_nsec,
+            "loss_stop": l_stop,
             "loss_type": l_type,
             "type_acc": type_acc,
             "loss_balance": l_balance,
             "loss_proc": l_proc,
             "loss_entropy": l_entropy,
             "nsec_acc": nsec_acc,
+            "stop_acc": stop_acc,
         }
 
     def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict:
@@ -742,10 +777,12 @@ class WGANStageTrainer(StageTrainer):
             "gp_loss",
             "loss_nsec",
             "nsec_acc",
+            "loss_stop",
+            "stop_acc",
             "grad_norm_d",
             "grad_norm_g",
         ]
-        if self.is_stage2 and self.particle_type_cfg.get("target") == "onehot":
+        if self.is_stage2 and self.particle_type_cfg.target == "onehot":
             # Differentiability instrumentation — only meaningful when the
             # type slice is a straight-through Gumbel relaxation.
             train_keys += ["grad_norm_type_slice", "grad_norm_cont_slice"]
@@ -754,11 +791,15 @@ class WGANStageTrainer(StageTrainer):
         self.stage_metrics = [stage_metric("lr"), stage_metric("critic_lr")]
 
     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
+        """Build `(real, fake_raw, mask, critic_fn, ar_inputs)` for stage 2,
+        covering both decoders and all three particle-type targets. `fake_raw`
+        still needs the caller's straight-through relaxation under
         `particle_type.target = "onehot"`, and neither tensor is masked-and-
-        multiplied on the fake side yet."""
+        multiplied on the fake side yet. `ar_inputs` is `None` under
+        `decoder = "one_shot"`; under `"autoregressive"` it's the same dict
+        `_ar_inputs` built to condition `self.model` above — returned so the
+        caller's `_stop_loss` reuses it instead of paying for a second
+        (possibly self-sampling) `_ar_inputs` call."""
         cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx = batch_tensors
         B = cond_cont.size(0)
         type_dim = stage2_type_dim(self.particle_type_cfg, self.particle_type_n_classes)
@@ -771,27 +812,28 @@ class WGANStageTrainer(StageTrainer):
         def critic_fn(x):
             return self.critic(x, cond_cont, cond_cat, stage1_ctx)
 
+        ar_inputs = None
         if self.decoder == "autoregressive":
             epoch = global_step // self.spec.steps_per_epoch
-            ar = self._ar_inputs(cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch)
-            real = self._sec_target(sec_cont, sec_type_idx, "wgan", flatten=False).reshape(B, -1) * mask
+            ar_inputs = self._ar_inputs(cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch)
+            real = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=False).reshape(B, -1) * mask
             z = torch.randn(B, k_max, self.model.noise_dim, device=device)
             fake_raw = self.model(
                 z,
                 cond_cont,
                 cond_cat,
                 stage1_ctx,
-                ar["history_feat"],
-                ar["has_prev"],
-                ar["remaining_frac"],
-                ar["slot_idx"],
+                ar_inputs["history_feat"],
+                ar_inputs["has_prev"],
+                ar_inputs["remaining_frac"],
+                ar_inputs["slot_idx"],
             ).reshape(B, -1)
         else:
-            real = self._sec_target(sec_cont, sec_type_idx, "wgan", flatten=True) * mask
+            real = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=True) * mask
             z = torch.randn(B, self.model.noise_dim, device=device)
             fake_raw = self.model(z, cond_cont, cond_cat, stage1_ctx)
 
-        return real, fake_raw, mask, critic_fn
+        return real, fake_raw, mask, critic_fn, ar_inputs
 
     def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict:
         (
@@ -807,6 +849,7 @@ class WGANStageTrainer(StageTrainer):
         stage1_ctx = x1_s1.detach()
         grad_probe: dict[str, float] = {}
 
+        ar_inputs = None
         if not self.is_stage2:
             real = x1_s1
 
@@ -817,13 +860,13 @@ class WGANStageTrainer(StageTrainer):
             fake = self.model(z, cond_cont, cond_cat)
             mask = None
         else:
-            real, fake_raw, mask, critic_fn = self._stage2_real_and_fake(
+            real, fake_raw, mask, critic_fn, ar_inputs = self._stage2_real_and_fake(
                 _Stage2RealFakeBatch(cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx),
                 stage1_ctx,
                 global_step,
                 device,
             )
-            if self.particle_type_cfg.get("target", "physical") == "onehot":
+            if self.particle_type_cfg.target == "onehot":
                 # Straight-through Gumbel-softmax relaxation of the type
                 # slice only — the critic must see a hard one-hot forward
                 # (matching what "real" data looks like) while gradient
@@ -859,18 +902,19 @@ class WGANStageTrainer(StageTrainer):
         # --- generator (+ n_sec) step ---
         did_g_step = global_step % self.n_critic == 0
         l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device)
+        l_stop, stop_acc = self._stop_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device, ar_inputs)
 
-        # On a non-generator-step batch with no n_sec_head on this stage
-        # (n_sec now defaults to stage 2), there's nothing for
+        # On a non-generator-step batch with no n_sec_head/stop_head on this
+        # stage (n_sec now defaults to stage 2), there's nothing for
         # the generator optimizer to do this batch — g_loss would otherwise
         # be a graph-less zero tensor, which .backward() rejects outright.
-        skip_g_step = not did_g_step and self.model.n_sec_head is None
+        skip_g_step = not did_g_step and self.model.n_sec_head is None and self.model.stop_head is None
         if did_g_step:
             g_loss_adv = generator_loss(critic_fn, fake)
-            g_loss = self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * l_nsec
+            g_loss = self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * (l_nsec + l_stop)
         else:
             g_loss_adv = torch.zeros((), device=device)
-            g_loss = self.spec.n_sec_lambda * l_nsec
+            g_loss = self.spec.n_sec_lambda * (l_nsec + l_stop)
         if skip_g_step:
             grad_norm_g = 0.0
         else:
@@ -888,6 +932,8 @@ class WGANStageTrainer(StageTrainer):
             "gp_loss": gp.item(),
             "loss_nsec": l_nsec.item(),
             "nsec_acc": nsec_acc.item(),
+            "loss_stop": l_stop.item(),
+            "stop_acc": stop_acc.item(),
             "did_g_step": did_g_step,
             "grad_norm": grad_norm_d + grad_norm_g,
             "grad_norm_d": grad_norm_d,
@@ -957,7 +1003,7 @@ def build_stage_trainers(
         if model is None:
             continue
         spec = StageSpec.from_config(cfg, name, is_stage2, max(total_train_batches, 1))
-        if spec.generator == "wgan":
+        if build_objective(spec.generator).is_adversarial:
             critic = critics.get(name)
             assert critic is not None, (
                 f"{name}_model.generator='wgan' requires a critic (see giant.model.network.build_critics)"

giant/validate.py

@@ -82,7 +82,7 @@ def validate_marginals(
     one-shot-vs-autoregressive-agnostic): n_sec
     distribution (+ classification accuracy), per-slot energy-fraction
     marginals, and a particle-type marginal whose shape depends on
-    `sec_decoder.particle_type_cfg["target"]` — restricted to each side's own
+    `sec_decoder.particle_type_cfg.target` — restricted to each side's own
     valid slots (real: `n_sec`; generated: the resolved `n_sec_pred`), since
     the two need not agree on how many slots are valid. Adds {"n_sec_real",
     "n_sec_pred", "n_sec_accuracy", "energy_fraction_kl"} plus, under
@@ -103,7 +103,7 @@ def validate_marginals(
         sec_decoder.eval()
 
     k_max = sec_decoder.k_max if sec_decoder is not None else 0
-    target = sec_decoder.particle_type_cfg.get("target", "physical") if sec_decoder is not None else "physical"
+    target = sec_decoder.particle_type_cfg.target if sec_decoder is not None else "physical"
 
     all_real, all_gen = [], []
     all_n_sec_real, all_n_sec_pred = [], []
@@ -129,10 +129,7 @@ def validate_marginals(
             continue
 
         n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cond_cont, cond_cat, gen, n_sec_pred)
-        n_sec_pred_np = n_sec_pred.cpu().numpy()
         n_sec_np = n_sec.numpy()
-        all_n_sec_real.append(n_sec_np)
-        all_n_sec_pred.append(n_sec_pred_np)
 
         real_valid = np.arange(k_max)[None, :] < n_sec_np[:, None]  # (B, k_max)
         real_frac = 1.0 / (1.0 + np.exp(-sec_cont[:, :, 0].numpy().astype(np.float64)))
@@ -140,6 +137,13 @@ def validate_marginals(
         sec_cont_pred, sec_type_pred, sec_valid_pred = sample_stage2(
             sec_decoder, cond_cont, cond_cat, gen, n_sec_pred, steps=steps
         )
+        # A stop-token decoder resolves n_sec_pred=None above — read the real
+        # count back off sec_valid_pred instead (a no-op round trip under
+        # every other n_sec.mode, where sec_valid_pred was built FROM
+        # n_sec_pred in the first place).
+        n_sec_pred_np = sec_valid_pred.sum(dim=-1).cpu().numpy()
+        all_n_sec_real.append(n_sec_np)
+        all_n_sec_pred.append(n_sec_pred_np)
         gen_frac = 1.0 / (1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64)))
         gen_valid = sec_valid_pred.cpu().numpy()
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "giant"
-version = "0.3.0"
+version = "0.3.2"
 description = "Geant4 step-function surrogate via conditional flow matching"
 readme = "README.md"
 requires-python = ">=3.12"

tests/test_analysis_reduce.py

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

tests/test_cond_layout.py

@@ -0,0 +1,87 @@
+import pytest
+from giant.cond_layout import AXIS_TYPES, CondLayout
+from giant.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM
+
+# ── cond_cat column layout ───────────────────────────────────────────────────
+
+
+def test_topn_cols_neither_onehot():
+    layout = CondLayout.from_types("physical", "embedding")
+    assert (layout.particle_topn_col, layout.material_topn_col) == (None, None)
+    assert layout.cat_dim == 2
+
+
+def test_topn_cols_particle_only():
+    layout = CondLayout.from_types("onehot", "physical")
+    assert (layout.particle_topn_col, layout.material_topn_col) == (2, None)
+    assert layout.cat_dim == 3
+
+
+def test_topn_cols_material_only():
+    layout = CondLayout.from_types("physical", "onehot")
+    assert (layout.particle_topn_col, layout.material_topn_col) == (None, 2)
+    assert layout.cat_dim == 3
+
+
+def test_topn_cols_both_onehot_particle_then_material():
+    layout = CondLayout.from_types("onehot", "onehot")
+    assert (layout.particle_topn_col, layout.material_topn_col) == (2, 3)
+    assert layout.cat_dim == 4
+
+
+def test_dense_vocab_cols_are_mode_independent():
+    """Columns 0/1 are always the dense pdg/material index — giant.model.routers
+    reads them without knowing the conditioning mode."""
+    assert (CondLayout.PDG_COL, CondLayout.MAT_COL) == (0, 1)
+    for particle in AXIS_TYPES:
+        for material in AXIS_TYPES:
+            layout = CondLayout.from_types(particle, material)
+            assert layout.particle_topn_col not in (layout.PDG_COL, layout.MAT_COL)
+            assert layout.material_topn_col not in (layout.PDG_COL, layout.MAT_COL)
+
+
+# ── cond_cont slice layout ───────────────────────────────────────────────────
+
+
+def test_cont_slices_tile_cond_cont_exactly():
+    """base / particle_phys / material_phys must partition cond_cont with no
+    gap and no overlap — a gap or overlap is exactly the silent
+    mis-indexing this object exists to prevent."""
+    layout = CondLayout.from_types("physical", "physical")
+    covered = list(range(*layout.base.indices(COND_DIM)))
+    covered += list(range(*layout.particle_phys.indices(COND_DIM)))
+    covered += list(range(*layout.material_phys.indices(COND_DIM)))
+    assert covered == list(range(COND_DIM))
+
+
+def test_cont_slice_widths_match_constants():
+    layout = CondLayout.from_types("embedding", "embedding")
+    assert layout.base == slice(0, COND_DIM_BASE)
+    assert layout.particle_phys.stop - layout.particle_phys.start == PARTICLE_PHYS_DIM
+    assert layout.material_phys.stop - layout.material_phys.start == MATERIAL_PHYS_DIM
+    assert layout.cont_dim == COND_DIM
+
+
+def test_cont_slices_are_mode_independent():
+    """cond_cont is COND_DIM wide in every mode — a non-"physical" axis gets
+    its block zero-filled rather than dropped, so the slices never move."""
+    physical = CondLayout.from_types("physical", "physical")
+    for particle in AXIS_TYPES:
+        for material in AXIS_TYPES:
+            layout = CondLayout.from_types(particle, material)
+            assert layout.base == physical.base
+            assert layout.particle_phys == physical.particle_phys
+            assert layout.material_phys == physical.material_phys
+
+
+# ── validation ───────────────────────────────────────────────────────────────
+
+
+def test_unknown_particle_type_raises():
+    with pytest.raises(ValueError, match="unknown conditioning.particle.type 'bogus'"):
+        CondLayout.from_types("bogus", "physical")
+
+
+def test_unknown_material_type_raises():
+    with pytest.raises(ValueError, match="unknown conditioning.material.type 'bogus'"):
+        CondLayout.from_types("physical", "bogus")

tests/test_config.py

@@ -51,9 +51,13 @@ def test_giant_config_to_dict_matches_default_config():
         gconfig.Stage2WganConfig,
         gconfig.RouterConfig,
         gconfig.Stage2RouterConfig,
+        gconfig.TrunkConfig,
         gconfig.NSecConfig,
         gconfig.ParticleTypeConfig,
         gconfig.AutoregressiveConfig,
+        gconfig.HeadConfig,
+        gconfig.Stage1HeadsConfig,
+        gconfig.Stage2HeadsConfig,
         gconfig.Stage1ModelConfig,
         gconfig.Stage2ModelConfig,
         gconfig.TrainConfig,
@@ -75,6 +79,39 @@ def test_stage2_model_config_defaults_match_documented_v030_intent():
     assert spec.particle_type.target == "onehot"
 
 
+def test_trunk_config_defaults_to_resmlp_for_both_stages():
+    """gitea #33: a v0.2-migrated / pre-existing config with no `trunk` key
+    at all must reproduce today's behaviour exactly."""
+    assert gconfig.Stage1ModelConfig().trunk.type == "resmlp"
+    assert gconfig.Stage2ModelConfig().trunk.type == "resmlp"
+    assert gconfig.DEFAULT_CONFIG["stage1_model"]["trunk"]["type"] == "resmlp"
+    assert gconfig.DEFAULT_CONFIG["stage2_model"]["trunk"]["type"] == "resmlp"
+
+
+def test_trunk_config_defaults_block_conditioning_to_add_for_both_stages():
+    """gitea #34: a pre-existing config with no `block_conditioning` key
+    must reproduce today's additive-bias behaviour exactly."""
+    assert gconfig.Stage1ModelConfig().trunk.block_conditioning == "add"
+    assert gconfig.Stage2ModelConfig().trunk.block_conditioning == "add"
+    assert gconfig.DEFAULT_CONFIG["stage1_model"]["trunk"]["block_conditioning"] == "add"
+    assert gconfig.DEFAULT_CONFIG["stage2_model"]["trunk"]["block_conditioning"] == "add"
+
+
+def test_heads_config_defaults_reproduce_pre_gitea_36_hardcoded_shape():
+    """gitea #36: a pre-existing config with no `heads` key must reproduce
+    today's hardcoded `hidden_dim // 2`, one-hidden-layer architecture
+    exactly."""
+    assert gconfig.Stage1ModelConfig().heads.n_sec.hidden_ratio == 0.5
+    assert gconfig.Stage1ModelConfig().heads.n_sec.depth == 2
+    assert gconfig.Stage2ModelConfig().heads.n_sec.hidden_ratio == 0.5
+    assert gconfig.Stage2ModelConfig().heads.n_sec.depth == 2
+    assert gconfig.Stage2ModelConfig().heads.type.hidden_ratio == 0.5
+    assert gconfig.Stage2ModelConfig().heads.type.depth == 2
+    assert gconfig.DEFAULT_CONFIG["stage1_model"]["heads"]["n_sec"] == {"hidden_ratio": 0.5, "depth": 2}
+    assert gconfig.DEFAULT_CONFIG["stage2_model"]["heads"]["n_sec"] == {"hidden_ratio": 0.5, "depth": 2}
+    assert gconfig.DEFAULT_CONFIG["stage2_model"]["heads"]["type"] == {"hidden_ratio": 0.5, "depth": 2}
+
+
 def test_particle_type_config_n_classes_defaults_to_zero_and_round_trips():
     """gitea #29: n_classes=0 means "inherit conditioning.particle.emb_dim"
     — the default must stay 0 so an existing config.toml with no
@@ -113,7 +150,17 @@ def test_n_sec_config_owner_defaults_to_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"}
+    assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1", "stop_sampling": "greedy"}
+
+
+def test_n_sec_config_stop_sampling_defaults_to_greedy():
+    assert gconfig.NSecConfig().stop_sampling == "greedy"
+
+
+def test_n_sec_config_stop_sampling_round_trips():
+    n_sec = gconfig.NSecConfig.from_dict({"mode": "stop_token", "stop_sampling": "sample"})
+    assert n_sec.stop_sampling == "sample"
+    assert n_sec.to_dict()["stop_sampling"] == "sample"
 
 
 # ---------------------------------------------------------------------------
@@ -688,13 +735,49 @@ def test_validate_config_tie_to_stage1_requires_stage1_active():
         assert "tie_to_stage1" in str(e)
 
 
-def test_validate_config_stop_token_not_implemented():
+def test_validate_config_stop_token_accepted_under_autoregressive():
+    """DEFAULT_CONFIG's stage2_model.decoder is already "autoregressive"
+    (see test_stage2_model_config_defaults_match_documented_v030_intent), so
+    mode="stop_token" alone must not raise."""
     cfg = _cfg_with(**{"stage2_model.n_sec.mode": "stop_token"})
+    gconfig.validate_config(cfg)  # must not raise
+
+
+def test_validate_config_stop_token_rejected_under_one_shot():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.n_sec.mode": "stop_token",
+            "stage2_model.decoder": "one_shot",
+        }
+    )
     try:
         gconfig.validate_config(cfg)
         assert False, "expected ValueError"
     except ValueError as e:
-        assert "stop_token" in str(e)
+        assert "stop_token" in str(e) and "autoregressive" in str(e)
+
+
+def test_validate_config_stop_token_rejected_for_stage1_owner():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.n_sec.mode": "stop_token",
+            "stage2_model.n_sec.owner": "stage1",
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "stop_token" in str(e) and "owner" in str(e)
+
+
+def test_validate_config_bad_stop_sampling_rejected():
+    cfg = _cfg_with(**{"stage2_model.n_sec.stop_sampling": "bogus"})
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "stop_sampling" in str(e)
 
 
 def test_validate_config_stage1_context_sampled_not_implemented():
@@ -911,6 +994,41 @@ def test_validate_config_keys_skips_meta_section():
     gconfig.validate_config_keys(cfg)  # must not raise
 
 
+def test_validate_config_keys_allows_trunk_type():
+    cfg = _cfg_with(**{"stage1_model.trunk.type": "resmlp", "stage2_model.trunk.type": "resmlp"})
+    gconfig.validate_config_keys(cfg)  # must not raise
+
+
+def test_validate_config_keys_rejects_unknown_trunk_key():
+    cfg = _cfg_with(**{"stage1_model.trunk.type_o": "resmlp"})  # typo for type
+    try:
+        gconfig.validate_config_keys(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "stage1_model.trunk.type_o" in str(e)
+        assert "type" in str(e)
+
+
+def test_validate_config_keys_allows_block_conditioning():
+    cfg = _cfg_with(
+        **{
+            "stage1_model.trunk.block_conditioning": "film",
+            "stage2_model.trunk.block_conditioning": "adaln",
+        }
+    )
+    gconfig.validate_config_keys(cfg)  # must not raise
+
+
+def test_validate_config_keys_rejects_unknown_block_conditioning_key():
+    cfg = _cfg_with(**{"stage1_model.trunk.block_conditioning_o": "film"})  # typo
+    try:
+        gconfig.validate_config_keys(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "stage1_model.trunk.block_conditioning_o" in str(e)
+        assert "block_conditioning" in str(e)
+
+
 def test_merge_cli_overrides_rejects_typo_in_toml_file(tmp_path):
     path = tmp_path / "config.toml"
     path.write_text("[meta]\nconfig_version = 3\n\n[stage1_model]\nn_res_block = 12\n")

tests/test_dwarf.py

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

tests/test_flow.py

@@ -1,11 +1,12 @@
 import torch
+from giant.config import ConditioningAxisConfig
 from giant.constants import COND_DIM
 from giant.model.network import Stage1Model
 from giant.model.schedule import CosineSchedule, flow_matching_loss
 from giant.sample import sample_flow, sample_ddim
 
-PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
-MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
+PARTICLE_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
+MATERIAL_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
 
 
 def _small_model():

tests/test_layers.py

@@ -0,0 +1,42 @@
+import pytest
+import torch
+
+from giant.model.layers import build_mlp_head
+
+
+def test_build_mlp_head_depth_1_is_bare_linear():
+    head = build_mlp_head(8, 4, hidden=16, depth=1)
+    assert len(head) == 1
+    assert isinstance(head[0], torch.nn.Linear)
+    assert head[0].in_features == 8
+    assert head[0].out_features == 4
+    out = head(torch.randn(3, 8))
+    assert out.shape == (3, 4)
+
+
+def test_build_mlp_head_depth_2_matches_pre_gitea_36_shape():
+    head = build_mlp_head(8, 4, hidden=16, depth=2)
+    assert len(head) == 3
+    assert isinstance(head[0], torch.nn.Linear)
+    assert head[0].in_features == 8
+    assert head[0].out_features == 16
+    assert isinstance(head[1], torch.nn.SiLU)
+    assert isinstance(head[2], torch.nn.Linear)
+    assert head[2].in_features == 16
+    assert head[2].out_features == 4
+    out = head(torch.randn(5, 8))
+    assert out.shape == (5, 4)
+
+
+def test_build_mlp_head_depth_3_has_extra_hidden_layer():
+    head = build_mlp_head(8, 4, hidden=16, depth=3)
+    assert len(head) == 5
+    widths = [(m.in_features, m.out_features) for m in head if isinstance(m, torch.nn.Linear)]
+    assert widths == [(8, 16), (16, 16), (16, 4)]
+    out = head(torch.randn(2, 8))
+    assert out.shape == (2, 4)
+
+
+def test_build_mlp_head_depth_0_raises():
+    with pytest.raises(ValueError, match="depth"):
+        build_mlp_head(8, 4, hidden=16, depth=0)

tests/test_network.py

@@ -5,24 +5,28 @@ import torch
 from giant import config as gconfig
 from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
 from giant.model.network import (
+    HISTORY_REGISTRY,
     AttentionHistory,
     ConditionEncoder,
+    HistoryEncoder,
     MarkovHistory,
     SinusoidalEmbedding,
     Stage1Model,
     Stage2Autoregressive,
     Stage2OneShot,
+    StageModel,
     build_critics,
+    build_history,
     build_models,
-    cat_col_layout,
+    build_objective,
     stage2_trunk_sec_dim,
     stage2_type_dim,
 )
 
-PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
-MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
-ONEHOT_PARTICLE_CFG = {"type": "onehot", "emb_dim": 6, "n_layers": 1}
-ONEHOT_MATERIAL_CFG = {"type": "onehot", "emb_dim": 4, "n_layers": 1}
+PARTICLE_CFG = gconfig.ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
+MATERIAL_CFG = gconfig.ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
+ONEHOT_PARTICLE_CFG = gconfig.ConditioningAxisConfig(type="onehot", emb_dim=6, n_layers=1)
+ONEHOT_MATERIAL_CFG = gconfig.ConditioningAxisConfig(type="onehot", emb_dim=4, n_layers=1)
 
 
 def test_sinusoidal_embedding_shape():
@@ -90,48 +94,72 @@ def test_stage1_model_no_n_sec_head_by_default():
     assert model.n_sec_head is None
 
 
-# --- cat_col_layout / stage2_type_dim / stage2_trunk_sec_dim ---------------
+def test_stage1_model_n_sec_head_default_cfg_matches_pre_gitea_36_shape():
+    """No n_sec_head_cfg given must reproduce the old hardcoded
+    hidden_dim // 2, one-hidden-layer architecture exactly (gitea #36)."""
+    model = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=40,
+        cond_out_dim=12,
+        n_sec_head_k_max=15,
+    )
+    assert model.n_sec_head is not None
+    assert len(model.n_sec_head) == 3
+    assert model.n_sec_head[0].in_features == 12
+    assert model.n_sec_head[0].out_features == 20  # hidden_dim // 2
+    assert model.n_sec_head[2].out_features == 16  # k_max + 1
 
 
-def test_cat_col_layout_neither_onehot():
-    assert cat_col_layout("physical", "embedding") == (None, None)
+def test_stage1_model_n_sec_head_cfg_controls_hidden_width_and_depth():
+    model = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=32,
+        cond_out_dim=16,
+        n_sec_head_k_max=15,
+        n_sec_head_cfg={"hidden_ratio": 0.25, "depth": 1},
+    )
+    assert model.n_sec_head is not None
+    assert len(model.n_sec_head) == 1
+    assert model.n_sec_head[0].in_features == 16
+    assert model.n_sec_head[0].out_features == 16
 
 
-def test_cat_col_layout_particle_only():
-    assert cat_col_layout("onehot", "physical") == (2, None)
-
-
-def test_cat_col_layout_material_only():
-    assert cat_col_layout("physical", "onehot") == (None, 2)
-
-
-def test_cat_col_layout_both_onehot_particle_then_material():
-    assert cat_col_layout("onehot", "onehot") == (2, 3)
+# --- stage2_type_dim / stage2_trunk_sec_dim --------------------------------
+# (the cond_cat column-layout tests live in tests/test_cond_layout.py)
 
 
 def test_stage2_type_dim_physical_is_particle_phys_dim():
-    assert stage2_type_dim({"target": "physical"}, emb_dim=16) == PARTICLE_PHYS_DIM
+    assert stage2_type_dim(gconfig.ParticleTypeConfig(target="physical"), emb_dim=16) == PARTICLE_PHYS_DIM
 
 
 def test_stage2_type_dim_onehot_and_embedding_are_emb_dim():
-    assert stage2_type_dim({"target": "onehot"}, emb_dim=16) == 16
-    assert stage2_type_dim({"target": "embedding"}, emb_dim=16) == 16
+    assert stage2_type_dim(gconfig.ParticleTypeConfig(target="onehot"), emb_dim=16) == 16
+    assert stage2_type_dim(gconfig.ParticleTypeConfig(target="embedding"), emb_dim=16) == 16
 
 
 def test_stage2_trunk_sec_dim_physical_matches_v02_sec_dim():
     k_max = 15
-    assert stage2_trunk_sec_dim({"target": "physical"}, "flow", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM
-    assert stage2_trunk_sec_dim({"target": "physical"}, "wgan", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM
+    physical = gconfig.ParticleTypeConfig(target="physical")
+    assert stage2_trunk_sec_dim(physical, "flow", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM
+    assert stage2_trunk_sec_dim(physical, "wgan", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM
 
 
 def test_stage2_trunk_sec_dim_onehot_wgan_folds_type_in():
     k_max = 15
-    assert stage2_trunk_sec_dim({"target": "onehot"}, "wgan", k_max, emb_dim=16) == k_max * (CONT_SLOT_DIM + 16)
+    onehot = gconfig.ParticleTypeConfig(target="onehot")
+    assert stage2_trunk_sec_dim(onehot, "wgan", k_max, emb_dim=16) == k_max * (CONT_SLOT_DIM + 16)
 
 
 def test_stage2_trunk_sec_dim_onehot_flow_excludes_type():
     k_max = 15
-    assert stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim=16) == k_max * CONT_SLOT_DIM
+    onehot = gconfig.ParticleTypeConfig(target="onehot")
+    assert stage2_trunk_sec_dim(onehot, "flow", k_max, emb_dim=16) == k_max * CONT_SLOT_DIM
 
 
 # --- ConditionEncoder onehot mode -------------------------------------------
@@ -139,8 +167,8 @@ def test_stage2_trunk_sec_dim_onehot_flow_excludes_type():
 
 def test_condition_encoder_onehot_forward_shape_and_gradients():
     B = 8
-    particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"])
-    material_emb_dim = int(ONEHOT_MATERIAL_CFG["emb_dim"])
+    particle_emb_dim = ONEHOT_PARTICLE_CFG.emb_dim
+    material_emb_dim = ONEHOT_MATERIAL_CFG.emb_dim
     enc = ConditionEncoder(
         pdg_vocab=5,
         mat_vocab=3,
@@ -171,12 +199,12 @@ def test_condition_encoder_onehot_is_a_true_one_hot_vector():
     verify the concatenated input segment really is one-hot, not e.g. an
     accidentally-learned embedding."""
     B = 4
-    particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"])
+    particle_emb_dim = ONEHOT_PARTICLE_CFG.emb_dim
     enc = ConditionEncoder(
         pdg_vocab=5,
         mat_vocab=3,
         particle_cfg=ONEHOT_PARTICLE_CFG,
-        material_cfg={"type": "physical", "emb_dim": 4, "n_layers": 1},
+        material_cfg=gconfig.ConditioningAxisConfig(type="physical", emb_dim=4, n_layers=1),
         out_dim=16,
     )
     cond_cont = torch.zeros(B, COND_DIM)
@@ -198,13 +226,11 @@ def test_condition_encoder_onehot_is_a_true_one_hot_vector():
 
 
 def _build_stage2(target: str, generator: str, emb_dim: int = 6) -> Stage2OneShot:
-    particle_cfg = {"type": "physical", "emb_dim": emb_dim, "n_layers": 1}
-    if target != "physical":
-        particle_cfg = dict(particle_cfg)
-        if target == "embedding":
-            particle_cfg["type"] = "embedding"
+    particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=emb_dim, n_layers=1)
+    if target == "embedding":
+        particle_cfg = gconfig.ConditioningAxisConfig(type="embedding", emb_dim=emb_dim, n_layers=1)
     k_max = 5
-    sec_dim = stage2_trunk_sec_dim({"target": target}, generator, k_max, emb_dim)
+    sec_dim = stage2_trunk_sec_dim(gconfig.ParticleTypeConfig(target=target), generator, k_max, emb_dim)
     return Stage2OneShot(
         pdg_vocab=5,
         mat_vocab=3,
@@ -217,7 +243,7 @@ def _build_stage2(target: str, generator: str, emb_dim: int = 6) -> Stage2OneSho
         sec_dim=sec_dim,
         generator=generator,
         k_max=k_max,
-        particle_type_cfg={"target": target, "lambda": 1.0},
+        particle_type_cfg=gconfig.ParticleTypeConfig(target=target),
     )
 
 
@@ -265,6 +291,38 @@ def test_stage2_oneshot_predict_type_raises_when_no_type_head():
         pass
 
 
+def test_stage2_oneshot_n_sec_head_and_type_head_cfg_control_hidden_width_and_depth():
+    """gitea #36: n_sec_head_cfg/type_head_cfg are independently tunable."""
+    k_max, emb_dim = 5, 6
+    particle_cfg = gconfig.ConditioningAxisConfig(type="onehot", emb_dim=emb_dim, n_layers=1)
+    sec_dim = stage2_trunk_sec_dim(gconfig.ParticleTypeConfig(target="onehot"), "flow", k_max, emb_dim)
+    model = Stage2OneShot(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=particle_cfg,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=40,
+        n_res_blocks=1,
+        cond_out_dim=12,
+        context_dim=8,
+        sec_dim=sec_dim,
+        generator="flow",
+        k_max=k_max,
+        particle_type_cfg=gconfig.ParticleTypeConfig(target="onehot"),
+        n_sec_head_cfg={"hidden_ratio": 0.25, "depth": 1},
+        type_head_cfg={"hidden_ratio": 0.75, "depth": 2},
+    )
+    assert model.n_sec_head is not None
+    assert len(model.n_sec_head) == 1
+    assert model.n_sec_head[0].in_features == 12
+    assert model.n_sec_head[0].out_features == k_max + 1
+
+    assert model.type_head is not None
+    assert len(model.type_head) == 3
+    assert model.type_head[0].out_features == 30  # round(40 * 0.75)
+    assert model.type_head[2].out_features == k_max * emb_dim
+
+
 def test_stage2_oneshot_forward_shape_onehot_wgan():
     B, k_max, emb_dim = 4, 5, 6
     model = _build_stage2("onehot", "wgan", emb_dim=emb_dim)
@@ -293,8 +351,8 @@ def test_stage2_oneshot_particle_type_n_classes_overrides_conditioning_emb_dim()
     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}
+    particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=6, n_layers=1)
+    particle_type_cfg = gconfig.ParticleTypeConfig(target="onehot", n_classes=20)
     sec_dim = stage2_trunk_sec_dim(particle_type_cfg, "flow", k_max, 20)
     model = Stage2OneShot(
         pdg_vocab=5,
@@ -310,7 +368,7 @@ def test_stage2_oneshot_particle_type_n_classes_overrides_conditioning_emb_dim()
         k_max=k_max,
         particle_type_cfg=particle_type_cfg,
     )
-    assert model.type_dim == 20  # not particle_cfg["emb_dim"] == 6
+    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
 
@@ -407,6 +465,48 @@ def test_attention_history_step_matches_forward():
     assert torch.allclose(stepped, expected, atol=1e-5)
 
 
+# --- HISTORY_REGISTRY / build_history (gitea #35) ----------------------------
+
+
+def test_history_registry_has_exactly_the_two_known_histories():
+    assert set(HISTORY_REGISTRY) == {"markov", "attention"}
+
+
+def test_build_history_returns_correct_concrete_type():
+    assert isinstance(build_history("markov", 4, 6), MarkovHistory)
+    assert isinstance(build_history("attention", 4, 8), AttentionHistory)
+
+
+def test_build_history_unknown_name_raises():
+    with pytest.raises(ValueError):
+        build_history("bogus", 4, 6)
+
+
+def test_build_history_filters_kwargs_by_signature():
+    """Attention-only kwargs (n_heads/n_layers) must be silently dropped when
+    building a MarkovHistory, matching build_router's documented behavior for
+    per-type hyperparameters coexisting in one config."""
+    hist = build_history("markov", 4, 6, n_heads=2, n_layers=1)
+    assert isinstance(hist, MarkovHistory)
+
+
+def test_history_encoder_base_default_init_cache_and_step():
+    """A HistoryEncoder subclass implementing only forward() must still get
+    working O(1) init_cache/step defaults from the base class."""
+
+    class _StubHistory(HistoryEncoder):
+        def forward(self, feat, has_prev):
+            return feat * 2
+
+    hist = _StubHistory()
+    assert hist.init_cache() is None
+    feat = torch.randn(2, 1, 4)
+    has_prev = torch.ones(2, 1, dtype=torch.bool)
+    out, cache = hist.step(feat, has_prev, "unused-cache")
+    assert torch.equal(out, hist.forward(feat, has_prev))
+    assert cache == "unused-cache"
+
+
 # --- Stage2Autoregressive (v0.3.0 step 5) -----------------------------------
 
 
@@ -417,10 +517,9 @@ def _build_stage2_ar(
     k_max: int = 5,
     history: str = "markov",
 ) -> Stage2Autoregressive:
-    particle_cfg = {"type": "physical", "emb_dim": emb_dim, "n_layers": 1}
+    particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=emb_dim, n_layers=1)
     if target == "embedding":
-        particle_cfg = dict(particle_cfg)
-        particle_cfg["type"] = "embedding"
+        particle_cfg = gconfig.ConditioningAxisConfig(type="embedding", emb_dim=emb_dim, n_layers=1)
     return Stage2Autoregressive(
         pdg_vocab=5,
         mat_vocab=3,
@@ -432,7 +531,7 @@ def _build_stage2_ar(
         context_dim=8,
         generator=generator,
         k_max=k_max,
-        particle_type_cfg={"target": target, "lambda": 1.0},
+        particle_type_cfg=gconfig.ParticleTypeConfig(target=target),
         history=history,
     )
 
@@ -453,8 +552,8 @@ def test_stage2_autoregressive_history_invalid_raises():
 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}
+    particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=6, n_layers=1)
+    particle_type_cfg = gconfig.ParticleTypeConfig(target="onehot", n_classes=20)
     model = Stage2Autoregressive(
         pdg_vocab=5,
         mat_vocab=3,
@@ -468,11 +567,42 @@ def test_stage2_autoregressive_particle_type_n_classes_overrides_conditioning_em
         k_max=5,
         particle_type_cfg=particle_type_cfg,
     )
-    assert model.type_dim == 20  # not particle_cfg["emb_dim"] == 6
+    assert model.type_dim == 20  # not particle_cfg.emb_dim == 6
     assert model.type_head is not None
     assert model.type_head[-1].out_features == 20
 
 
+def test_stage2_autoregressive_n_sec_head_and_type_head_cfg_control_hidden_width_and_depth():
+    """gitea #36, Stage2Autoregressive side — see the Stage2OneShot version
+    of this test for the full rationale."""
+    particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=6, n_layers=1)
+    particle_type_cfg = gconfig.ParticleTypeConfig(target="onehot")
+    model = Stage2Autoregressive(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=particle_cfg,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=40,
+        n_res_blocks=1,
+        cond_out_dim=12,
+        context_dim=8,
+        generator="flow",
+        k_max=5,
+        particle_type_cfg=particle_type_cfg,
+        n_sec_head_cfg={"hidden_ratio": 0.25, "depth": 1},
+        type_head_cfg={"hidden_ratio": 0.75, "depth": 2},
+    )
+    assert model.n_sec_head is not None
+    assert len(model.n_sec_head) == 1
+    assert model.n_sec_head[0].in_features == 12
+    assert model.n_sec_head[0].out_features == 6  # k_max + 1
+
+    assert model.type_head is not None
+    assert len(model.type_head) == 3
+    assert model.type_head[0].out_features == 30  # round(40 * 0.75)
+    assert model.type_head[2].out_features == model.type_dim
+
+
 @pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
 @pytest.mark.parametrize("generator", ["wgan", "flow"])
 @pytest.mark.parametrize("history", ["markov", "attention"])
@@ -482,9 +612,9 @@ def test_stage2_autoregressive_forward_shape(target, generator, history):
     cond_cont = torch.randn(B, COND_DIM)
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     stage1_out = torch.randn(B, 9)
-    type_dim = stage2_type_dim({"target": target}, emb_dim)
+    type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target=target), emb_dim)
     history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
-    token_dim = stage2_trunk_sec_dim({"target": target}, generator, 1, emb_dim)
+    token_dim = stage2_trunk_sec_dim(gconfig.ParticleTypeConfig(target=target), generator, 1, emb_dim)
     if generator == "wgan":
         x_t = torch.randn(B, K, model.noise_dim)
         t = None
@@ -521,7 +651,7 @@ def test_stage2_autoregressive_predict_type_shape():
     cond_cont = torch.randn(B, COND_DIM)
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     stage1_out = torch.randn(B, 9)
-    type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
+    type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target="onehot"), emb_dim)
     history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
     out = model.predict_type(
         cond_cont,
@@ -542,7 +672,7 @@ def test_stage2_autoregressive_predict_type_raises_when_no_type_head(target, gen
     cond_cont = torch.randn(B, COND_DIM)
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     stage1_out = torch.randn(B, 9)
-    type_dim = stage2_type_dim({"target": target}, emb_dim)
+    type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target=target), emb_dim)
     history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
     with pytest.raises(RuntimeError):
         model.predict_type(
@@ -562,7 +692,7 @@ def test_stage2_autoregressive_gradients_flow_wgan_onehot():
     cond_cont = torch.randn(B, COND_DIM)
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     stage1_out = torch.randn(B, 9)
-    type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
+    type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target="onehot"), emb_dim)
     history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
     z = torch.randn(B, K, model.noise_dim)
     gen_out = model(
@@ -587,9 +717,9 @@ def test_stage2_autoregressive_gradients_flow_onehot():
     cond_cont = torch.randn(B, COND_DIM)
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     stage1_out = torch.randn(B, 9)
-    type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
+    type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target="onehot"), emb_dim)
     history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
-    token_dim = stage2_trunk_sec_dim({"target": "onehot"}, "flow", 1, emb_dim)
+    token_dim = stage2_trunk_sec_dim(gconfig.ParticleTypeConfig(target="onehot"), "flow", 1, emb_dim)
     x_t = torch.randn(B, K, token_dim)
     t = torch.rand(B, K)
     flow_out = model(
@@ -629,7 +759,7 @@ def test_stage2_autoregressive_history_step_matches_parallel_history_encoder():
     B, K, emb_dim = 3, 6, 6
     model = _build_stage2_ar("physical", "wgan", emb_dim=emb_dim, k_max=K, history="attention")
     model.eval()
-    type_dim = stage2_type_dim({"target": "physical"}, emb_dim)
+    type_dim = stage2_type_dim(gconfig.ParticleTypeConfig(target="physical"), emb_dim)
     hist_in_dim = CONT_SLOT_DIM + type_dim
     own_feat = torch.randn(B, K, hist_in_dim)  # token i's own raw feature
     has_prev_full = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
@@ -705,6 +835,54 @@ def test_build_models_share_stages_true_shared_params_are_in_both_stage_paramete
     assert shared_ids <= {id(p) for p in stage2.parameters()}
 
 
+def test_condition_encoder_stores_the_exact_particle_and_material_cfg_instances_passed_in():
+    """gitea #38: ConditionEncoder must not round-trip particle_cfg/
+    material_cfg through a dict — the exact ConditioningAxisConfig instance
+    passed in is what `.particle_cfg`/`.material_cfg` hold afterward."""
+    particle_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
+    material_cfg = gconfig.ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
+    enc = ConditionEncoder(pdg_vocab=3, mat_vocab=2, particle_cfg=particle_cfg, material_cfg=material_cfg)
+    assert enc.particle_cfg is particle_cfg
+    assert enc.material_cfg is material_cfg
+
+
+def test_stagemodel_stores_the_exact_particle_type_cfg_instance_passed_in():
+    """gitea #38: a StageModel subclass must not round-trip particle_type_cfg
+    through a dict — the exact ParticleTypeConfig instance passed in is what
+    `.particle_type_cfg` holds afterward."""
+    particle_type_cfg = gconfig.ParticleTypeConfig(target="onehot", n_classes=11)
+    sec_dim = stage2_trunk_sec_dim(particle_type_cfg, "flow", 5, 11)
+    model = Stage2OneShot(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        k_max=5,
+        sec_dim=sec_dim,
+        particle_type_cfg=particle_type_cfg,
+    )
+    assert model.particle_type_cfg is particle_type_cfg
+
+
+def test_build_models_particle_type_cfg_and_conditioning_axes_are_dataclasses_not_dicts():
+    """gitea #38: build_models must pass the parsed ConditioningAxisConfig/
+    ParticleTypeConfig dataclasses themselves down to the model constructors,
+    not re-serialize them to a dict first (the inversion the issue names) —
+    before the fix, .particle_type_cfg was a plain dict (s2_spec.particle_type
+    .to_dict()) and .cond_enc.particle_cfg came from the raw, unparsed
+    conditioning["particle"] dict."""
+    cfg = _minimal_model_config(share_stages=False)
+    cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0, "n_classes": 11}
+    built = build_models(cfg)
+    stage1, stage2 = built["stage1"], built["stage2"]
+    assert stage1 is not None and stage2 is not None
+    assert isinstance(stage2.particle_type_cfg, gconfig.ParticleTypeConfig)
+    assert isinstance(stage1.cond_enc.particle_cfg, gconfig.ConditioningAxisConfig)
+    assert isinstance(stage1.cond_enc.material_cfg, gconfig.ConditioningAxisConfig)
+
+
 def test_build_models_particle_type_n_classes_overrides_conditioning_emb_dim():
     """gitea #29 end-to-end through build_models: setting
     stage2_model.particle_type.n_classes independently of
@@ -765,7 +943,40 @@ def _partial_model_config() -> dict:
 def test_build_models_omitted_decoder_and_particle_type_match_default_config():
     built = build_models(_partial_model_config())
     assert isinstance(built["stage2"], Stage2Autoregressive)
-    assert built["stage2"].particle_type_cfg["target"] == "onehot"
+    assert built["stage2"].particle_type_cfg.target == "onehot"
+
+
+def test_build_models_custom_heads_block_controls_head_shapes():
+    """gitea #36: stage{1,2}_model.heads flows all the way from config dict
+    through build_models to the actual constructed head shapes."""
+    cfg = _partial_model_config()
+    cfg["stage1_model"] = {
+        "active": True,
+        "hidden_dim": 40,
+        "n_res_blocks": 1,
+        "heads": {"n_sec": {"hidden_ratio": 0.25, "depth": 1}},
+    }
+    cfg["stage2_model"]["decoder"] = "one_shot"
+    cfg["stage2_model"]["generator"] = "flow"  # wgan folds the type slice; no separate type_head
+    cfg["stage2_model"]["n_sec"] = {"owner": "stage1"}
+    cfg["stage2_model"]["heads"] = {
+        "n_sec": {"hidden_ratio": 0.25, "depth": 1},
+        "type": {"hidden_ratio": 0.75, "depth": 2},
+    }
+
+    built = build_models(cfg)
+    stage1, stage2 = built["stage1"], built["stage2"]
+    assert stage1 is not None
+    assert stage2 is not None
+
+    assert stage1.n_sec_head is not None
+    assert len(stage1.n_sec_head) == 1  # owner=stage1, so stage1 builds it
+    assert stage2.n_sec_head is None  # owner=stage1, so stage2 doesn't
+
+    assert isinstance(stage2, Stage2OneShot)
+    assert stage2.type_head is not None
+    assert len(stage2.type_head) == 3
+    assert stage2.type_head[0].out_features == 6  # round(8 * 0.75)
 
 
 def test_build_critics_omitted_particle_type_matches_default_config():
@@ -826,3 +1037,163 @@ def test_build_critics_stage2_critic_hidden_dim_and_n_res_blocks_override_genera
     assert overridden is not None
     assert overridden.input_proj.out_features == 16
     assert len(overridden.blocks) == 3
+
+
+# ── StageModel base (gitea #39): Stage1Model/Stage2OneShot/Stage2Autoregressive
+# scaffolding — construction order, and therefore fresh-init RNG draw order and
+# state_dict key set, must stay byte-for-byte what it was before the base class
+# existed. ------------------------------------------------------------------
+
+_STAGE_HIDDEN_DIM = 32
+_STAGE_N_BLOCKS = 2
+_STAGE_COND_OUT_DIM = 16
+
+
+def _resblock_keys(prefix: str) -> set[str]:
+    return {
+        f"{prefix}.norm.weight",
+        f"{prefix}.norm.bias",
+        f"{prefix}.linear1.weight",
+        f"{prefix}.linear1.bias",
+        f"{prefix}.cond_proj.weight",
+        f"{prefix}.linear2.weight",
+        f"{prefix}.linear2.bias",
+    }
+
+
+def _trunk_keys(prefix: str = "trunk") -> set[str]:
+    keys = {
+        f"{prefix}.input_proj.weight",
+        f"{prefix}.input_proj.bias",
+        f"{prefix}.out_proj.weight",
+        f"{prefix}.out_proj.bias",
+    }
+    for i in range(_STAGE_N_BLOCKS):
+        keys |= _resblock_keys(f"{prefix}.blocks.{i}")
+    return keys
+
+
+def _cond_enc_keys() -> set[str]:
+    return {
+        "cond_enc.mlp.0.weight",
+        "cond_enc.mlp.0.bias",
+        "cond_enc.mlp.2.weight",
+        "cond_enc.mlp.2.bias",
+        "cond_enc.particle_mlp.0.weight",
+        "cond_enc.particle_mlp.0.bias",
+        "cond_enc.material_mlp.0.weight",
+        "cond_enc.material_mlp.0.bias",
+    }
+
+
+def _fuse_keys(name: str) -> set[str]:
+    return {f"{name}.0.weight", f"{name}.0.bias"}
+
+
+def _head_keys(name: str) -> set[str]:
+    return {f"{name}.0.weight", f"{name}.0.bias", f"{name}.2.weight", f"{name}.2.bias"}
+
+
+def _expected_stage_keys(*, has_time: bool, extra: set[str]) -> set[str]:
+    keys = _cond_enc_keys() | _trunk_keys() | extra
+    if has_time:
+        keys.add("time_emb.freqs")
+    return keys
+
+
+@pytest.mark.parametrize("generator", ["flow", "wgan"])
+def test_stage1_model_state_dict_keys_unchanged_by_stagemodel_refactor(generator):
+    model = Stage1Model(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=_STAGE_HIDDEN_DIM,
+        n_res_blocks=_STAGE_N_BLOCKS,
+        cond_out_dim=_STAGE_COND_OUT_DIM,
+        generator=generator,
+        time_dim=8,
+        noise_dim=8,
+        n_sec_head_k_max=15,
+    )
+    expected = _expected_stage_keys(
+        has_time=build_objective(generator).needs_time,
+        extra=_head_keys("n_sec_head"),
+    )
+    assert set(model.state_dict().keys()) == expected
+
+
+@pytest.mark.parametrize("generator", ["flow", "wgan"])
+def test_stage2_oneshot_state_dict_keys_unchanged_by_stagemodel_refactor(generator):
+    model = Stage2OneShot(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=_STAGE_HIDDEN_DIM,
+        n_res_blocks=_STAGE_N_BLOCKS,
+        cond_out_dim=_STAGE_COND_OUT_DIM,
+        generator=generator,
+        time_dim=8,
+        noise_dim=8,
+        k_max=15,
+    )
+    extra = _head_keys("n_sec_head") | {"context_adapter.proj.weight", "context_adapter.proj.bias"} | _fuse_keys("fuse")
+    expected = _expected_stage_keys(has_time=build_objective(generator).needs_time, extra=extra)
+    assert set(model.state_dict().keys()) == expected
+
+
+@pytest.mark.parametrize("generator", ["flow", "wgan"])
+def test_stage2_autoregressive_state_dict_keys_unchanged_by_stagemodel_refactor(generator):
+    model = Stage2Autoregressive(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=_STAGE_HIDDEN_DIM,
+        n_res_blocks=_STAGE_N_BLOCKS,
+        cond_out_dim=_STAGE_COND_OUT_DIM,
+        generator=generator,
+        time_dim=8,
+        noise_dim=8,
+        k_max=15,
+    )
+    extra = (
+        _head_keys("n_sec_head")
+        | {"context_adapter.proj.weight", "context_adapter.proj.bias"}
+        | _fuse_keys("base_fuse")
+        | _fuse_keys("token_fuse")
+        | {"history_encoder.start", "history_encoder.mlp.0.weight", "history_encoder.mlp.0.bias"}
+    )
+    expected = _expected_stage_keys(has_time=build_objective(generator).needs_time, extra=extra)
+    assert set(model.state_dict().keys()) == expected
+
+
+@pytest.mark.parametrize("cls", [Stage1Model, Stage2OneShot, Stage2Autoregressive])
+def test_stage_classes_are_stagemodel_subclasses(cls):
+    assert issubclass(cls, StageModel)
+
+
+@pytest.mark.parametrize("cls", [Stage1Model, Stage2OneShot, Stage2Autoregressive])
+@pytest.mark.parametrize("generator", ["flow", "ddpm", "wgan"])
+def test_stagemodel_time_emb_matches_objective_needs_time(cls, generator):
+    kwargs = dict(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=_STAGE_HIDDEN_DIM,
+        n_res_blocks=_STAGE_N_BLOCKS,
+        cond_out_dim=_STAGE_COND_OUT_DIM,
+        generator=generator,
+        time_dim=8,
+        noise_dim=8,
+    )
+    if cls is Stage1Model:
+        kwargs["n_sec_head_k_max"] = 15
+    else:
+        kwargs["k_max"] = 15
+    model = cls(**kwargs)
+    assert model.generator_kind == generator
+    assert model.noise_dim == 8
+    assert (model.time_emb is not None) == build_objective(generator).needs_time

tests/test_objectives.py

@@ -0,0 +1,251 @@
+"""Tests for `giant/model/objectives.py` — the generator/objective registry
+(gitea #32) that replaced bare `generator in ("flow", "ddpm", "wgan")`
+string checks scattered across models.py/sample.py/builders.py/
+stage2_inputs.py/trainers.py."""
+
+import pytest
+import torch
+
+from giant.config import ConditioningAxisConfig
+from giant.constants import COND_DIM, CONT_SLOT_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM
+from giant.model.network import (
+    OBJECTIVE_REGISTRY,
+    DdpmObjective,
+    FlowObjective,
+    Stage1Model,
+    Stage2Autoregressive,
+    Stage2OneShot,
+    WganObjective,
+    build_objective,
+)
+from giant.model.schedule import CosineSchedule, flow_matching_loss, flow_matching_loss_secondary
+
+_PHYS_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
+
+
+def _cond(B: int, pdg: int = 3, mat: int = 2) -> tuple[torch.Tensor, torch.Tensor]:
+    cond_cont = torch.randn(B, COND_DIM)
+    cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
+    return cond_cont, cond_cat
+
+
+# ── registry ─────────────────────────────────────────────────────────────
+
+
+def test_registry_has_exactly_the_three_known_objectives():
+    assert set(OBJECTIVE_REGISTRY) == {"flow", "ddpm", "wgan"}
+
+
+def test_build_objective_returns_correct_concrete_type():
+    assert isinstance(build_objective("flow"), FlowObjective)
+    assert isinstance(build_objective("ddpm"), DdpmObjective)
+    assert isinstance(build_objective("wgan"), WganObjective)
+
+
+def test_build_objective_unknown_name_raises():
+    with pytest.raises(ValueError, match="unknown generator/objective"):
+        build_objective("bogus")
+
+
+def test_build_objective_filters_kwargs_by_signature():
+    # FlowObjective takes no constructor args — n_steps (a DdpmObjective-only
+    # kwarg) must be silently dropped, not raise a TypeError.
+    build_objective("flow", n_steps=500)
+    ddpm = build_objective("ddpm", n_steps=250)
+    assert isinstance(ddpm, DdpmObjective)
+    assert ddpm.n_steps == 250
+
+
+# ── flags ────────────────────────────────────────────────────────────────
+
+
+def test_flow_objective_flags():
+    obj = build_objective("flow")
+    assert obj.needs_time is True
+    assert obj.is_adversarial is False
+    assert obj.folds_type_slice is False
+    assert obj.supports_stage2_decoder is True
+
+
+def test_ddpm_objective_flags():
+    obj = build_objective("ddpm")
+    assert obj.needs_time is True
+    assert obj.is_adversarial is False
+    assert obj.folds_type_slice is False
+    assert obj.supports_stage2_decoder is False
+
+
+def test_wgan_objective_flags():
+    obj = build_objective("wgan")
+    assert obj.needs_time is False
+    assert obj.is_adversarial is True
+    assert obj.folds_type_slice is True
+    assert obj.supports_stage2_decoder is True
+
+
+# ── trunk_in_dim ─────────────────────────────────────────────────────────
+
+
+def test_trunk_in_dim_flow_and_ddpm_pass_through_out_dim():
+    assert build_objective("flow").trunk_in_dim(out_dim=9, noise_dim=8) == 9
+    assert build_objective("ddpm").trunk_in_dim(out_dim=9, noise_dim=8) == 9
+
+
+def test_trunk_in_dim_wgan_uses_noise_dim():
+    assert build_objective("wgan").trunk_in_dim(out_dim=9, noise_dim=8) == 8
+
+
+# ── ddpm schedule ────────────────────────────────────────────────────────
+
+
+def test_ddpm_build_schedule_has_requested_length():
+    schedule = build_objective("ddpm").build_schedule(n_steps=17, device=torch.device("cpu"))
+    assert isinstance(schedule, CosineSchedule)
+    assert schedule.T == 17
+
+
+def test_flow_and_wgan_build_schedule_is_none():
+    assert build_objective("flow").build_schedule(100, torch.device("cpu")) is None
+    assert build_objective("wgan").build_schedule(100, torch.device("cpu")) is None
+
+
+# ── stage1_loss parity ──────────────────────────────────────────────────
+
+
+def test_flow_objective_stage1_loss_matches_direct_call():
+    torch.manual_seed(0)
+    model = Stage1Model(
+        pdg_vocab=3, mat_vocab=2, particle_cfg=_PHYS_CFG, material_cfg=_PHYS_CFG, hidden_dim=16, n_res_blocks=1
+    )
+    cond_cont, cond_cat = _cond(4)
+    x1 = torch.randn(4, X_DIM)
+
+    torch.manual_seed(1)
+    expected = flow_matching_loss(model, x1, cond_cont, cond_cat)
+    torch.manual_seed(1)
+    actual = build_objective("flow").stage1_loss(model, x1, cond_cont, cond_cat)
+    assert torch.allclose(actual, expected)
+
+
+def test_ddpm_objective_stage1_loss_matches_direct_call():
+    torch.manual_seed(0)
+    model = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=_PHYS_CFG,
+        material_cfg=_PHYS_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        generator="ddpm",
+    )
+    cond_cont, cond_cat = _cond(4)
+    x1 = torch.randn(4, X_DIM)
+    objective = build_objective("ddpm", n_steps=50)
+    schedule = objective.build_schedule(50, torch.device("cpu"))
+    assert isinstance(schedule, CosineSchedule)
+
+    torch.manual_seed(1)
+    expected = schedule.loss(model, x1, cond_cont, cond_cat)
+    torch.manual_seed(1)
+    actual = objective.stage1_loss(model, x1, cond_cont, cond_cat, schedule=schedule)
+    assert torch.allclose(actual, expected)
+
+
+def test_ddpm_objective_stage1_loss_requires_a_schedule():
+    model = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=_PHYS_CFG,
+        material_cfg=_PHYS_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        generator="ddpm",
+    )
+    cond_cont, cond_cat = _cond(4)
+    with pytest.raises(AssertionError):
+        build_objective("ddpm").stage1_loss(model, torch.randn(4, X_DIM), cond_cont, cond_cat, schedule=None)
+
+
+# ── stage2_loss dispatch ─────────────────────────────────────────────────
+
+
+def test_flow_objective_stage2_loss_one_shot_matches_direct_call():
+    torch.manual_seed(0)
+    B, k_max = 4, 5
+    sec_dim = k_max * SEC_SLOT_DIM
+    model = Stage2OneShot(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=_PHYS_CFG,
+        material_cfg=_PHYS_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        generator="flow",
+        sec_dim=sec_dim,
+        k_max=k_max,
+    )
+    cond_cont, cond_cat = _cond(B)
+    stage1_ctx = torch.randn(B, X_DIM)
+    x1_s2 = torch.randn(B, sec_dim)
+    sec_mask = torch.ones(B, k_max, dtype=torch.bool)
+
+    torch.manual_seed(1)
+    expected = flow_matching_loss_secondary(model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=None)
+    torch.manual_seed(1)
+    actual = build_objective("flow").stage2_loss(
+        model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=None, ar_inputs=None
+    )
+    assert torch.allclose(actual, expected)
+
+
+def test_flow_objective_stage2_loss_dispatches_to_ar_when_ar_inputs_given():
+    torch.manual_seed(0)
+    B, k_max = 4, 5
+    model = Stage2Autoregressive(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=_PHYS_CFG,
+        material_cfg=_PHYS_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        generator="flow",
+        k_max=k_max,
+    )
+    cond_cont, cond_cat = _cond(B)
+    stage1_ctx = torch.randn(B, X_DIM)
+    token_dim = CONT_SLOT_DIM + PARTICLE_PHYS_DIM
+    x1_s2 = torch.randn(B, k_max, token_dim)
+    sec_mask = torch.ones(B, k_max, dtype=torch.bool)
+    ar_inputs = {
+        "history_feat": torch.randn(B, k_max, token_dim),
+        "has_prev": torch.ones(B, k_max, dtype=torch.bool),
+        "remaining_frac": torch.rand(B, k_max),
+        "slot_idx": torch.linspace(0, 1, k_max).unsqueeze(0).expand(B, -1),
+    }
+
+    loss = build_objective("flow").stage2_loss(
+        model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=None, ar_inputs=ar_inputs
+    )
+    assert loss.dim() == 0
+    assert torch.isfinite(loss)
+
+
+def test_ddpm_objective_stage2_loss_not_implemented():
+    dummy_model = torch.nn.Module()
+    dummy = torch.zeros(1)
+    with pytest.raises(NotImplementedError):
+        build_objective("ddpm").stage2_loss(
+            dummy_model, dummy, dummy, dummy, dummy, torch.ones(1, 1, dtype=torch.bool), type_dim=None
+        )
+
+
+def test_wgan_objective_has_no_loss_methods():
+    dummy_model = torch.nn.Module()
+    dummy = torch.zeros(1)
+    objective = build_objective("wgan")
+    with pytest.raises(NotImplementedError):
+        objective.stage1_loss(dummy_model, dummy, dummy, dummy)
+    with pytest.raises(NotImplementedError):
+        objective.stage2_loss(
+            dummy_model, dummy, dummy, dummy, dummy, torch.ones(1, 1, dtype=torch.bool), type_dim=None
+        )

tests/test_phase2.py

@@ -4,6 +4,7 @@ import numpy as np
 import pytest
 import torch
 
+from giant.config import ConditioningAxisConfig
 from giant.constants import (
     COND_DIM,
     CONT_SLOT_DIM,
@@ -23,9 +24,9 @@ from giant.sample import sample_secondaries
 # ── helpers ──────────────────────────────────────────────────────────────────
 
 
-def _particle_material_cfg(conditioning: str) -> tuple[dict, dict]:
-    cfg = {"type": conditioning, "emb_dim": 16, "n_layers": 1}
-    return dict(cfg), dict(cfg)
+def _particle_material_cfg(conditioning: str) -> tuple[ConditioningAxisConfig, ConditioningAxisConfig]:
+    cfg = ConditioningAxisConfig(type=conditioning, emb_dim=16, n_layers=1)
+    return cfg, cfg
 
 
 def _stage1(pdg=3, mat=2, conditioning="embedding"):
@@ -626,3 +627,53 @@ def test_decode_secondaries_mass_charge_round_trip_with_normalizer():
     _, _, sec_mass, sec_charge, _ = decode_secondaries(sec_cont_normed, n_sec, e_sec, pre_dir, sec_phys_normalizer=norm)
     assert sec_mass[0, 0] == pytest.approx(938.27208943, abs=1e-2)
     assert sec_charge[0, 0] == pytest.approx(1.0, abs=1e-4)
+
+
+def test_decode_secondaries_extreme_negative_log_mass_stays_nonnegative():
+    from giant.data.transforms import decode_secondaries, log_transform
+
+    N = 1
+    e_sec = np.array([5.0], dtype=np.float32)
+    n_sec = np.array([1])
+    pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
+
+    sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32)
+    sec_cont[0, 0, 0] = 10.0  # stick logit -> ~all of e_sec
+    sec_cont[0, 0, 1:4] = [0, 0, 1]
+    sec_cont[0, 0, 4] = -50.0  # raw model prediction: extremely negative log_mass
+    sec_cont[0, 0, 5] = 1.0
+
+    _, _, sec_mass, _, _ = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
+
+    # A raw model prediction isn't itself the output of log_transform, so
+    # naively applying inv_log_transform can undershoot zero (see
+    # decode_secondaries) — which then crashes the next log_transform call
+    # once this mass is fed back in as conditioning during rollout. The
+    # float32 residual from clipping can land a hair below zero, but must
+    # stay well above -eps so log_transform(mass) stays finite.
+    assert sec_mass[0, 0] > -1e-8
+    log_transform(sec_mass[0, 0])
+
+
+def test_decode_secondaries_extreme_positive_log_mass_stays_finite():
+    from giant.data.transforms import decode_secondaries, log_transform
+
+    N = 1
+    e_sec = np.array([5.0], dtype=np.float32)
+    n_sec = np.array([1])
+    pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
+
+    sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32)
+    sec_cont[0, 0, 0] = 10.0  # stick logit -> ~all of e_sec
+    sec_cont[0, 0, 1:4] = [0, 0, 1]
+    sec_cont[0, 0, 4] = 200.0  # raw model prediction: extremely positive log_mass
+    sec_cont[0, 0, 5] = 1.0
+
+    _, _, sec_mass, _, _ = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
+
+    # Mirror image of the extreme-negative case above: exp(log_mass)
+    # overflows float32 to inf for an unclipped raw prediction this large,
+    # which then crashes the next log_transform call the same way a
+    # negative mass would.
+    assert np.isfinite(sec_mass[0, 0])
+    log_transform(sec_mass[0, 0])

tests/test_rollout.py

@@ -8,6 +8,7 @@ import numpy as np
 import pytest
 import torch
 
+from giant.config import ConditioningAxisConfig, ParticleTypeConfig
 from giant.constants import TERM_ESCAPED, TERM_MAX_STEPS, TERM_UNKNOWN_PDG, K_MAX
 from giant.data.loader import TopNMap
 from giant.data.transforms import Normalizer
@@ -27,8 +28,8 @@ MAT_MAP = {"G4_AIR": 0, "G4_PbWO4": 1}
 
 
 def _models(conditioning="embedding"):
-    particle_cfg = {"type": conditioning, "emb_dim": 16, "n_layers": 1}
-    material_cfg = {"type": conditioning, "emb_dim": 16, "n_layers": 1}
+    particle_cfg = ConditioningAxisConfig(type=conditioning, emb_dim=16, n_layers=1)
+    material_cfg = ConditioningAxisConfig(type=conditioning, emb_dim=16, n_layers=1)
     s1 = Stage1Model(
         pdg_vocab=3,
         mat_vocab=2,
@@ -365,9 +366,10 @@ def _models_v3(
     k_max=6,
     emb_dim=4,
     stage2_has_n_sec_head=True,
+    stop_token=False,
 ):
-    particle_cfg = {"type": conditioning, "emb_dim": emb_dim, "n_layers": 1}
-    material_cfg = {"type": conditioning, "emb_dim": emb_dim, "n_layers": 1}
+    particle_cfg = ConditioningAxisConfig(type=conditioning, emb_dim=emb_dim, n_layers=1)
+    material_cfg = ConditioningAxisConfig(type=conditioning, emb_dim=emb_dim, n_layers=1)
     # A fresh v0.3.0 Stage1Model — no n_sec_head_k_max, unlike _models() above
     # (n_sec ownership moves to stage 2 by default).
     s1 = Stage1Model(
@@ -380,7 +382,7 @@ def _models_v3(
         generator=generator1,
         noise_dim=8,
     )
-    particle_type_cfg = {"target": target}
+    particle_type_cfg = ParticleTypeConfig(target=target)
     # Explicit kwargs rather than a shared **common dict: a dict() call whose
     # values have heterogeneous types (str/int/dict/bool) widens under static
     # analysis to dict[str, <big union>], which then makes every constructor
@@ -416,7 +418,8 @@ def _models_v3(
             noise_dim=8,
             k_max=k_max,
             particle_type_cfg=particle_type_cfg,
-            build_n_sec_head=stage2_has_n_sec_head,
+            build_n_sec_head=stage2_has_n_sec_head and not stop_token,
+            build_stop_head=stop_token,
         )
     return s1.eval(), s2.eval()
 
@@ -480,6 +483,23 @@ def test_rollout_stage2_owns_n_sec_when_stage1_has_no_head(fake_material_props):
         assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4)
 
 
+@pytest.mark.parametrize("generator2", ["flow", "wgan"])
+def test_rollout_stop_token_end_to_end(fake_material_props, generator2):
+    """gitea #40: n_sec.mode='stop_token' (no n_sec_head on either stage —
+    resolve_n_sec must return None and let sample_stage2's AR loop derive
+    the count from its own stop head) must still run to completion, produce
+    secondaries, and conserve energy exactly like the 'head' mode."""
+    s1, s2 = _models_v3(decoder="autoregressive", generator2=generator2, stop_token=True)
+    rec = _run_v3(s1, s2)
+    assert len(rec["event_id"]) > 0
+    seeds = _seeds()
+    for i, ev in enumerate(seeds["event_id"]):
+        m = rec["event_id"] == ev
+        dep = rec["edep"][m].sum()
+        leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum()
+        assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4)
+
+
 def test_resolve_n_sec_raises_when_neither_stage_owns_head(fake_material_props):
     """Neither stage owning n_sec_head only happens for a
     stage2_model.n_sec.mode other than "head" — not a valid rollout-capable
@@ -557,8 +577,8 @@ COND_MAT_TOPN_MAP = TopNMap(class_map={"G4_AIR": 0, "G4_PbWO4": 1}, other_member
 
 
 def _onehot_conditioning_models():
-    particle_cfg = {"type": "onehot", "emb_dim": len(PDG_MAP), "n_layers": 1}
-    material_cfg = {"type": "onehot", "emb_dim": len(MAT_MAP), "n_layers": 1}
+    particle_cfg = ConditioningAxisConfig(type="onehot", emb_dim=len(PDG_MAP), n_layers=1)
+    material_cfg = ConditioningAxisConfig(type="onehot", emb_dim=len(MAT_MAP), n_layers=1)
     s1 = Stage1Model(
         pdg_vocab=3,
         mat_vocab=2,
@@ -574,7 +594,7 @@ def _onehot_conditioning_models():
         material_cfg=material_cfg,
         hidden_dim=32,
         n_res_blocks=2,
-        sec_dim=stage2_trunk_sec_dim({"target": "physical"}, "flow", K_MAX, 3),
+        sec_dim=stage2_trunk_sec_dim(ParticleTypeConfig(target="physical"), "flow", K_MAX, 3),
         generator="flow",
         time_dim=16,
     )
@@ -637,9 +657,9 @@ def _run_conditioning_and_type_onehot_different_n_classes():
     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}
+    particle_cfg = ConditioningAxisConfig(type="onehot", emb_dim=cond_emb_dim, n_layers=1)
+    material_cfg = ConditioningAxisConfig(type="onehot", emb_dim=len(MAT_MAP), n_layers=1)
+    particle_type_cfg = ParticleTypeConfig(target="onehot", n_classes=type_n_classes)
 
     s1 = Stage1Model(
         pdg_vocab=3,

tests/test_router.py

@@ -3,24 +3,32 @@
 import pytest
 import torch
 
+from giant.config import ConditioningAxisConfig
 from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
 from giant.model.network import (
+    BLOCK_REGISTRY,
+    TRUNK_REGISTRY,
+    AdaLNResBlock,
     ComposedRouter,
     EnergyRouter,
-    MonolithicTrunk,
+    ExpertTrunk,
+    FilmResBlock,
     PdgRouter,
     ProcessRouter,
     ROUTER_REGISTRY,
+    ResBlock,
     RoutedTrunk,
     Stage1Model,
     Stage2OneShot,
+    build_block,
     build_composed_router,
+    build_expert_body,
     build_models,
     build_router,
 )
 
-PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
-MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
+PARTICLE_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
+MATERIAL_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
 
 
 def _cond(B=8, pdg=3, mat=2):
@@ -144,6 +152,68 @@ def test_build_router_unknown_type_raises():
     raise AssertionError("expected ValueError for unknown router type")
 
 
+# ── TRUNK_REGISTRY / build_expert_body ──────────────────────────────────────
+
+
+def test_trunk_registry_has_resmlp():
+    assert "resmlp" in TRUNK_REGISTRY
+    assert TRUNK_REGISTRY["resmlp"] is ExpertTrunk
+
+
+def test_build_expert_body_unknown_type_raises():
+    try:
+        build_expert_body("nonexistent", in_dim=4, out_dim=4, hidden_dim=8, n_blocks=1, cond_dim=4)
+    except ValueError:
+        return
+    raise AssertionError("expected ValueError for unknown trunk type")
+
+
+# ── BLOCK_REGISTRY / build_block (gitea #34) ────────────────────────────────
+
+
+def test_block_registry_has_add_film_adaln():
+    assert BLOCK_REGISTRY["add"] is ResBlock
+    assert BLOCK_REGISTRY["film"] is FilmResBlock
+    assert BLOCK_REGISTRY["adaln"] is AdaLNResBlock
+
+
+def test_build_block_unknown_type_raises():
+    try:
+        build_block("nonexistent", dim=8, cond_dim=4)
+    except ValueError:
+        return
+    raise AssertionError("expected ValueError for unknown block conditioning type")
+
+
+@pytest.mark.parametrize("block_type", ["add", "film", "adaln"])
+def test_block_forward_shape(block_type):
+    block = build_block(block_type, dim=8, cond_dim=4)
+    x = torch.randn(5, 8)
+    cond = torch.randn(5, 4)
+    out = block(x, cond)
+    assert out.shape == (5, 8)
+
+
+def test_film_res_block_output_invariant_to_cond_at_init():
+    """Zero-initialized film_proj means gamma=beta=0 at construction, so the
+    output must not depend on which cond is passed in."""
+    block = FilmResBlock(dim=8, cond_dim=4)
+    x = torch.randn(5, 8)
+    cond_a = torch.randn(5, 4)
+    cond_b = torch.randn(5, 4)
+    torch.testing.assert_close(block(x, cond_a), block(x, cond_b))
+
+
+def test_adaln_res_block_is_identity_at_init():
+    """Zero-initialized adaln_proj means scale=shift=gate=0 at construction,
+    so the block must be the exact identity function (the 'Zero' in
+    AdaLN-Zero)."""
+    block = AdaLNResBlock(dim=8, cond_dim=4)
+    x = torch.randn(5, 8)
+    cond = torch.randn(5, 4)
+    torch.testing.assert_close(block(x, cond), x)
+
+
 # ── EnergyRouter learn_width / learn_temperature ────────────────────────────
 
 
@@ -993,8 +1063,10 @@ def test_build_models_monolith_when_router_absent():
     stage1, stage2 = models["stage1"], models["stage2"]
     assert isinstance(stage1, Stage1Model)
     assert isinstance(stage2, Stage2OneShot)
-    assert isinstance(stage1.trunk, MonolithicTrunk)
-    assert isinstance(stage2.trunk, MonolithicTrunk)
+    assert not isinstance(stage1.trunk, RoutedTrunk)
+    assert not isinstance(stage2.trunk, RoutedTrunk)
+    assert isinstance(stage1.trunk, ExpertTrunk)
+    assert isinstance(stage2.trunk, ExpertTrunk)
 
 
 def test_build_models_monolith_when_router_disabled():
@@ -1006,8 +1078,10 @@ def test_build_models_monolith_when_router_disabled():
     models = build_models(cfg)
     stage1, stage2 = models["stage1"], models["stage2"]
     assert stage1 is not None and stage2 is not None
-    assert isinstance(stage1.trunk, MonolithicTrunk)
-    assert isinstance(stage2.trunk, MonolithicTrunk)
+    assert not isinstance(stage1.trunk, RoutedTrunk)
+    assert not isinstance(stage2.trunk, RoutedTrunk)
+    assert isinstance(stage1.trunk, ExpertTrunk)
+    assert isinstance(stage2.trunk, ExpertTrunk)
 
 
 def test_build_models_routed_when_enabled():
@@ -1040,6 +1114,59 @@ def test_build_models_routed_when_enabled():
     assert len(stage2.trunk.experts) == 4
 
 
+def test_build_models_explicit_resmlp_trunk_type_matches_default():
+    """stage1_model.trunk.type = 'resmlp' is the default's spelled-out
+    equivalent, not a behaviour change — gitea #33."""
+    default_cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2)
+    explicit_cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2, trunk={"type": "resmlp"})
+    default_stage1 = build_models(default_cfg)["stage1"]
+    explicit_stage1 = build_models(explicit_cfg)["stage1"]
+    assert default_stage1 is not None and explicit_stage1 is not None
+    assert type(default_stage1.trunk) is type(explicit_stage1.trunk) is ExpertTrunk
+    assert default_stage1.trunk.input_proj.weight.shape == explicit_stage1.trunk.input_proj.weight.shape
+    default_params = sum(p.numel() for p in default_stage1.parameters())
+    explicit_params = sum(p.numel() for p in explicit_stage1.parameters())
+    assert default_params == explicit_params
+
+
+def test_build_models_explicit_add_block_conditioning_matches_default():
+    """stage1_model.trunk.block_conditioning = 'add' is the default's
+    spelled-out equivalent, not a behaviour change — gitea #34."""
+    default_cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2)
+    explicit_cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2, trunk={"type": "resmlp", "block_conditioning": "add"})
+    default_stage1 = build_models(default_cfg)["stage1"]
+    explicit_stage1 = build_models(explicit_cfg)["stage1"]
+    assert default_stage1 is not None and explicit_stage1 is not None
+    assert type(default_stage1.trunk.blocks[0]) is type(explicit_stage1.trunk.blocks[0]) is ResBlock
+    default_params = sum(p.numel() for p in default_stage1.parameters())
+    explicit_params = sum(p.numel() for p in explicit_stage1.parameters())
+    assert default_params == explicit_params
+
+
+@pytest.mark.parametrize("block_type,cls", [("film", FilmResBlock), ("adaln", AdaLNResBlock)])
+def test_build_models_selects_block_conditioning(block_type, cls):
+    cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2, trunk={"type": "resmlp", "block_conditioning": block_type})
+    stage1 = build_models(cfg)["stage1"]
+    assert stage1 is not None
+    assert isinstance(stage1.trunk, ExpertTrunk)
+    assert all(isinstance(b, cls) for b in stage1.trunk.blocks)
+
+
+def test_build_models_routed_trunk_uses_block_conditioning_for_every_expert():
+    cfg = _nested_cfg(
+        pdg_vocab=4,
+        mat_vocab=2,
+        trunk={"type": "resmlp", "block_conditioning": "film"},
+        stage1_router={"enabled": True, "type": "energy", "n_experts": 3},
+    )
+    stage1 = build_models(cfg)["stage1"]
+    assert stage1 is not None
+    assert isinstance(stage1.trunk, RoutedTrunk)
+    assert len(stage1.trunk.experts) == 3
+    for expert in stage1.trunk.experts:
+        assert all(isinstance(b, FilmResBlock) for b in expert.blocks)
+
+
 def test_build_models_routed_pair_is_drop_in_for_sample_flow():
     """Exercise the exact calling convention giant/sample.py uses."""
     from giant.sample import sample_flow, sample_secondaries

tests/test_sample.py

@@ -5,6 +5,7 @@ for the one-shot samplers."""
 import pytest
 import torch
 
+from giant.config import ConditioningAxisConfig, ParticleTypeConfig
 from giant.constants import COND_DIM, CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, X_DIM
 from giant.model.network import (
     Stage1Model,
@@ -20,12 +21,12 @@ from giant.sample import (
     sample_wgan,
 )
 
-_PHYS_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
+_PHYS_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
 
 
-def _particle_material_cfg(conditioning: str, emb_dim: int) -> tuple[dict, dict]:
-    cfg = {"type": conditioning, "emb_dim": emb_dim, "n_layers": 1}
-    return dict(cfg), dict(cfg)
+def _particle_material_cfg(conditioning: str, emb_dim: int) -> tuple[ConditioningAxisConfig, ConditioningAxisConfig]:
+    cfg = ConditioningAxisConfig(type=conditioning, emb_dim=emb_dim, n_layers=1)
+    return cfg, cfg
 
 
 def _cond(B: int, pdg: int = 3, mat: int = 2) -> tuple[torch.Tensor, torch.Tensor]:
@@ -43,7 +44,7 @@ def _conditioning_for(target: str) -> str:
 
 def _stage2_oneshot(target: str, generator: str, emb_dim: int = 6, pdg: int = 3, mat: int = 2) -> Stage2OneShot:
     particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
-    particle_type_cfg = {"target": target}
+    particle_type_cfg = ParticleTypeConfig(target=target)
     # build_models (giant/model/network.py) computes sec_dim this same way
     # before constructing Stage2OneShot — its own default (SEC_DIM, the
     # "physical" width) is only correct for target="physical".
@@ -84,7 +85,7 @@ def _stage2_ar(
         time_dim=16,
         noise_dim=8,
         k_max=k_max,
-        particle_type_cfg={"target": target},
+        particle_type_cfg=ParticleTypeConfig(target=target),
         history=history,
         attn_n_heads=2,
         attn_n_layers=1,
@@ -95,6 +96,45 @@ def _expected_type_dim(target: str, emb_dim: int) -> int:
     return PARTICLE_PHYS_DIM if target == "physical" else emb_dim
 
 
+def _stage2_ar_stop_token(
+    target: str,
+    generator: str,
+    stop_sampling: str = "greedy",
+    emb_dim: int = 6,
+    pdg: int = 3,
+    mat: int = 2,
+    k_max: int = 5,
+) -> Stage2Autoregressive:
+    particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
+    return Stage2Autoregressive(
+        pdg_vocab=pdg,
+        mat_vocab=mat,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+        generator=generator,
+        time_dim=16,
+        noise_dim=8,
+        k_max=k_max,
+        particle_type_cfg=ParticleTypeConfig(target=target),
+        build_n_sec_head=False,
+        build_stop_head=True,
+        stop_sampling=stop_sampling,
+    ).eval()
+
+
+def _force_stop_head_logit(decoder: Stage2Autoregressive, logit: float) -> None:
+    """Zeroes stop_head's weights and pins its bias, so predict_stop returns
+    `logit` for every row/slot regardless of conditioning — makes the AR
+    loop's stop decision deterministic for testing."""
+    assert decoder.stop_head is not None
+    last_linear = decoder.stop_head[-1]
+    with torch.no_grad():
+        last_linear.weight.zero_()
+        last_linear.bias.fill_(logit)
+
+
 # ── Stage-1 n_sec ownership ──────────────────────────────────────────────────
 
 
@@ -222,3 +262,77 @@ def test_sample_secondaries_ar_first_slot_has_no_history():
     sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
     assert sec_cont.shape == (B, 1, CONT_SLOT_DIM)
     assert sec_valid.tolist() == [[False], [True], [True]]
+
+
+# ── Stage2Autoregressive: n_sec.mode = "stop_token" ─────────────────────────
+
+
+@pytest.mark.parametrize("stop_sampling", ["greedy", "sample"])
+def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(stop_sampling):
+    """A stop_head pinned to a large positive logit fires at slot 0 for
+    every row under both policies (greedy: sigmoid(logit) >= 0.5; sample:
+    a Bernoulli draw at sigmoid(logit) ~= 1) — the loop should break before
+    generating any token."""
+    B, k_max = 4, 5
+    decoder = _stage2_ar_stop_token("physical", "flow", stop_sampling=stop_sampling, k_max=k_max)
+    _force_stop_head_logit(decoder, 50.0)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
+    assert sec_valid.shape == (B, k_max)
+    assert not sec_valid.any()
+
+
+@pytest.mark.parametrize("stop_sampling", ["greedy", "sample"])
+def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(stop_sampling):
+    """A stop_head pinned to a large negative logit never fires under either
+    policy, so every row is capped at k_max (the safety cap, not a modeling
+    ceiling)."""
+    B, k_max = 4, 5
+    decoder = _stage2_ar_stop_token("physical", "flow", stop_sampling=stop_sampling, k_max=k_max)
+    _force_stop_head_logit(decoder, -50.0)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
+    assert sec_valid.all()
+
+
+@pytest.mark.parametrize("generator", ["flow", "wgan"])
+def test_sample_secondaries_ar_stop_token_valid_mask_is_always_a_prefix(generator):
+    """Without forcing the stop head, per-row stop timing varies — but
+    sec_valid must always be a contiguous prefix (slot k valid implies every
+    slot < k is also valid), matching the "head"/"truth" contract."""
+    B, k_max = 6, 5
+    decoder = _stage2_ar_stop_token("physical", generator, k_max=k_max)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    _, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
+    n = sec_valid.sum(dim=-1)
+    expected = torch.arange(k_max).unsqueeze(0) < n.unsqueeze(1)
+    assert torch.equal(sec_valid, expected)
+
+
+def test_sample_secondaries_ar_stop_token_explicit_n_sec_pred_ignores_stop_head():
+    """The scheduled-sampling training contract: passing n_sec_pred
+    explicitly (as _assemble_stage2_ar_inputs_scheduled's self-sample call
+    does, with ground-truth n_sec) must run the full k_max loop and mask by
+    the given count, even though the decoder owns a stop_head that would
+    otherwise stop early."""
+    B, k_max = 3, 5
+    decoder = _stage2_ar_stop_token("physical", "flow", k_max=k_max)
+    _force_stop_head_logit(decoder, 50.0)  # would stop immediately if consulted
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec_pred = torch.tensor([0, 2, k_max])
+    _, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
+    for i, n in enumerate(n_sec_pred.tolist()):
+        assert sec_valid[i, :n].all()
+        assert not sec_valid[i, n:].any()
+
+
+def test_sample_secondaries_ar_none_n_sec_pred_without_stop_head_raises():
+    decoder = _stage2_ar("physical", "flow", k_max=5)  # head mode: no stop_head
+    cond_cont, cond_cat = _cond(3)
+    stage1_out = torch.randn(3, X_DIM)
+    with pytest.raises(AssertionError):
+        sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)

tests/test_train.py

@@ -10,6 +10,7 @@ from unittest.mock import MagicMock
 import pytest
 import torch
 
+from giant.config import ParticleTypeConfig
 from giant.constants import (
     COND_DIM,
     CONT_SLOT_DIM,
@@ -19,7 +20,7 @@ from giant.constants import (
     X_DIM,
 )
 from giant.data.dataset import StepBatch
-from giant.model.network import build_critics, build_models
+from giant.model.network import Stage2Autoregressive, build_critics, build_models
 from giant.training import (
     FlowDDPMStageTrainer,
     StageSpec,
@@ -39,6 +40,7 @@ from giant.training.stage2_inputs import (
     _shift_prev,
     _stage2_tf_prob,
     _stick_fraction,
+    _stop_target_and_mask,
     _type_repr,
 )
 
@@ -119,6 +121,23 @@ def test_ar_has_prev_false_only_at_slot_zero():
     assert has_prev.tolist() == [[False, True, True, True, True]]
 
 
+def test_stop_target_and_mask_hand_computed():
+    # k_max=5; n_sec=0 (no real secondaries, stop slot is 0), n_sec=2
+    # (stop slot is 2), n_sec=5 (== k_max: no in-range stop slot at all).
+    n_sec = torch.tensor([0, 2, 5])
+    target, mask = _stop_target_and_mask(n_sec, 5, torch.device("cpu"))
+    assert target.tolist() == [
+        [1, 0, 0, 0, 0],
+        [0, 0, 1, 0, 0],
+        [0, 0, 0, 0, 0],
+    ]
+    assert mask.tolist() == [
+        [True, False, False, False, False],
+        [True, True, True, False, False],
+        [True, True, True, True, True],
+    ]
+
+
 # --- _stage2_tf_prob (v0.3.0 step 7) -----------
 
 
@@ -158,7 +177,7 @@ def test_type_repr_shapes_and_values(target):
     cond_enc = torch.nn.Module()
     if target == "embedding":
         cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim)
-    repr_ = _type_repr(sec_type_idx, sec_cont, {"target": target}, cond_enc, emb_dim)
+    repr_ = _type_repr(sec_type_idx, sec_cont, ParticleTypeConfig(target=target), cond_enc, emb_dim)
     expected_width = PARTICLE_PHYS_DIM if target == "physical" else emb_dim
     assert repr_.shape == (B, K, expected_width)
     if target == "physical":
@@ -187,7 +206,7 @@ def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened(target
     cond_enc = torch.nn.Module()
     if target == "embedding":
         cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim)
-    particle_type_cfg = {"target": target}
+    particle_type_cfg = ParticleTypeConfig(target=target)
     flat = _assemble_stage2_real(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim)
     unflat = _assemble_stage2_ar_target(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim)
     assert torch.equal(unflat.flatten(1), flat)
@@ -198,7 +217,7 @@ def test_assemble_stage2_ar_inputs_shapes_and_history_feat_width():
     sec_cont = torch.randn(B, K_MAX, SEC_SLOT_DIM)
     sec_type_idx = torch.randint(0, emb_dim, (B, K_MAX))
     cond_enc = torch.nn.Module()
-    out = _assemble_stage2_ar_inputs(sec_cont, sec_type_idx, {"target": "physical"}, cond_enc, emb_dim)
+    out = _assemble_stage2_ar_inputs(sec_cont, sec_type_idx, ParticleTypeConfig(target="physical"), cond_enc, emb_dim)
     assert out["history_feat"].shape == (B, K_MAX, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
     assert out["has_prev"].shape == (B, K_MAX)
     assert out["remaining_frac"].shape == (B, K_MAX)
@@ -739,6 +758,71 @@ def test_wgan_onehot_one_shot_also_gets_grad_norm_instrumentation():
         assert any(float(r["stage2/train/grad_norm_cont_slice"]) > 0 for r in rows)
 
 
+# --- n_sec.mode = "stop_token" (gitea #40) ----------------------------------
+
+
+def _stop_token_cfg():
+    cfg = _base_cfg()
+    cfg["stage2_model"]["decoder"] = "autoregressive"
+    cfg["stage2_model"]["n_sec"] = {"mode": "stop_token", "lambda": 0.1}
+    return cfg
+
+
+@pytest.mark.parametrize("stage2_generator", ["wgan", "flow"])
+def test_build_stage_trainers_stop_token_step_runs(stage2_generator):
+    """A stop_token AR stage-2 trainer.step() must run and emit a finite
+    loss_stop for both non-adversarial (flow) and WGAN generators — the two
+    trainer subclasses wire the stop head's BCE term in independently."""
+    cfg = _stop_token_cfg()
+    cfg["stage2_model"]["generator"] = stage2_generator
+    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)
+    trainer = trainers["stage2"]
+    batch = _fake_batches(1, 4)[0]
+    stats = trainer.step(batch, torch.device("cpu"), global_step=1)
+    assert math.isfinite(stats["loss_stop"])
+    assert math.isfinite(stats["stop_acc"])
+
+
+def test_stop_token_model_has_stop_head_not_n_sec_head():
+    cfg = _stop_token_cfg()
+    model_config = _model_config(cfg)
+    models = build_models(model_config)
+    stage2 = models["stage2"]
+    assert isinstance(stage2, Stage2Autoregressive)
+    assert stage2.n_sec_head is None
+    assert stage2.stop_head is not None
+
+
+def test_head_mode_model_has_n_sec_head_not_stop_head():
+    """Sanity check on the other side of the gate — the default 'head' mode
+    must be unaffected by the stop_head plumbing."""
+    cfg = _base_cfg()
+    cfg["stage2_model"]["decoder"] = "autoregressive"
+    model_config = _model_config(cfg)
+    models = build_models(model_config)
+    stage2 = models["stage2"]
+    assert isinstance(stage2, Stage2Autoregressive)
+    assert stage2.n_sec_head is not None
+    assert stage2.stop_head is None
+
+
+def test_train_end_to_end_stop_token():
+    """Full train() run with n_sec.mode='stop_token' must complete and write
+    a checkpoint + metrics.csv with finite losses throughout."""
+    cfg = _stop_token_cfg()
+    with tempfile.TemporaryDirectory() as tmp:
+        out_dir = Path(tmp) / "run"
+        _run_train(cfg, out_dir)
+        assert (out_dir / "last.pt").exists()
+        with open(out_dir / "metrics.csv", newline="") as f:
+            rows = list(csv.DictReader(f))
+        assert len(rows) == cfg["train"]["epochs"]
+        assert all(math.isfinite(float(r["stage2/train/loss_stop"])) for r in rows)
+
+
 def test_wgan_physical_omits_grad_norm_slice_columns():
     cfg = _base_cfg()  # _base_cfg's stage2_model.particle_type.target is "physical"
     with tempfile.TemporaryDirectory() as tmp:

tests/test_transforms.py

@@ -2,6 +2,7 @@ import warnings
 
 import numpy as np
 import pytest
+from giant.cond_layout import AXIS_TYPES, CondLayout
 from giant.constants import COND_DIM, COND_DIM_BASE, K_MAX
 from giant.data.transforms import (
     build_cond_features,
@@ -531,6 +532,129 @@ def test_build_cond_features_rejects_legacy_normalizer_in_physical_mode(
         )
 
 
+# ── build_cond_features / build_features share one column layout (gitea #37) ──
+
+
+@pytest.mark.parametrize("particle_type", AXIS_TYPES)
+@pytest.mark.parametrize("material_type", AXIS_TYPES)
+def test_both_builders_agree_column_for_column(particle_type, material_type, fake_material_props):
+    """The two builders used to lay out cond_cont/cond_cat independently and
+    drift apart silently. They now share `_build_cond_arrays`, so for every
+    mode pair they must produce identical arrays."""
+    data = _minimal_step_data(3)
+    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
+    pdg_topn = {11: 0} if particle_type == "onehot" else None
+    mat_topn = {"PbWO4": 0} if material_type == "onehot" else None
+
+    cond_cont, cond_cat = build_cond_features(
+        data,
+        pdg_map,
+        mat_map,
+        particle_conditioning=particle_type,
+        material_conditioning=material_type,
+        pdg_topn_map=pdg_topn,
+        mat_topn_map=mat_topn,
+    )
+    feats = build_features(
+        data,
+        pdg_map,
+        mat_map,
+        particle_conditioning=particle_type,
+        material_conditioning=material_type,
+        pdg_topn_map=pdg_topn,
+        mat_topn_map=mat_topn,
+    )
+
+    layout = CondLayout.from_types(particle_type, material_type)
+    assert cond_cat.shape[1] == layout.cat_dim
+    np.testing.assert_array_equal(feats.cond_cont, cond_cont)
+    np.testing.assert_array_equal(feats.cond_cat, cond_cat)
+
+
+def test_build_features_physical_mode_tolerates_out_of_vocab_pdg_and_material():
+    """The permissive vocab lookup added for "physical"/"onehot" mode (see
+    build_cond_features) applies to build_features too — `giant predict` on a
+    file whose pdg/material aren't in the checkpoint's dense vocab must not
+    KeyError when nothing reads those indices."""
+    pdg_map = {11: 0, 22: 1}
+    mat_map = {"G4_AIR": 0}
+    data = _minimal_step_data(2)
+    data["pdg"] = np.full(2, 13, dtype=np.int64)  # not in pdg_map
+    data["material"] = np.full(2, "G4_Pb", dtype=object)  # not in mat_map
+
+    _, cond_cat, *_ = build_features(
+        data,
+        pdg_map,
+        mat_map,
+        particle_conditioning="physical",
+        material_conditioning="physical",
+    )
+    np.testing.assert_array_equal(cond_cat, [[0, 0], [0, 0]])  # dummy indices, no raise
+
+    with pytest.raises(KeyError):
+        build_features(
+            data,
+            pdg_map,
+            mat_map,
+            particle_conditioning="embedding",
+            material_conditioning="embedding",
+        )
+
+
+def test_build_features_pads_legacy_normalizer_in_embedding_mode():
+    """The legacy-normalizer padding (a pre-physical-conditioning checkpoint's
+    cond normalizer is COND_DIM_BASE wide) applies to build_features too —
+    `giant predict` reaches build_features, not build_cond_features."""
+    data = _minimal_step_data(3)
+    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
+    legacy_norm = Normalizer()
+    legacy_norm.mean = np.zeros(COND_DIM_BASE, dtype=np.float32)
+    legacy_norm.std = np.ones(COND_DIM_BASE, dtype=np.float32)
+
+    cond_cont, *_ = build_features(
+        data,
+        pdg_map,
+        mat_map,
+        cond_normalizer=legacy_norm,
+        particle_conditioning="embedding",
+        material_conditioning="embedding",
+    )
+
+    assert cond_cont.shape[-1] == COND_DIM
+    np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE:], 0.0)
+
+
+def test_build_features_rejects_legacy_normalizer_in_physical_mode(fake_material_props):
+    data = _minimal_step_data(3)
+    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
+    legacy_norm = Normalizer()
+    legacy_norm.mean = np.zeros(COND_DIM_BASE, dtype=np.float32)
+    legacy_norm.std = np.ones(COND_DIM_BASE, dtype=np.float32)
+
+    with pytest.raises(ValueError, match="predates physical-property conditioning"):
+        build_features(
+            data,
+            pdg_map,
+            mat_map,
+            cond_normalizer=legacy_norm,
+            particle_conditioning="physical",
+            material_conditioning="physical",
+        )
+
+
+def test_onehot_axis_without_its_topn_map_raises():
+    """`cond_cat`'s width is the layout's call, so a "onehot" axis with no
+    top-N map is a hard error rather than a silently-narrower array that
+    ConditionEncoder would then index out of bounds."""
+    data = _minimal_step_data(2)
+    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
+
+    with pytest.raises(ValueError, match="needs pdg_topn_map"):
+        build_cond_features(data, pdg_map, mat_map, particle_conditioning="onehot")
+    with pytest.raises(ValueError, match="needs mat_topn_map"):
+        build_cond_features(data, pdg_map, mat_map, material_conditioning="onehot")
+
+
 # ── sorted_membership / _vectorized_map_lookup ──────────────────────────────
 
 

tests/test_validate.py

@@ -1,17 +1,18 @@
 import numpy as np
 import torch
 
+from giant.config import ConditioningAxisConfig, ParticleTypeConfig
 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
 
-_PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
-_MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
+_PARTICLE_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
+_MATERIAL_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
 _K_MAX = 5
 
 
-def _tiny_models(particle_type_cfg: dict | None = None):
+def _tiny_models(particle_type_cfg: ParticleTypeConfig | None = None):
     """A fresh v0.3.0 pair: Stage1Model owns no n_sec_head, so n_sec always
     comes from Stage2OneShot."""
     s1 = Stage1Model(
@@ -22,12 +23,13 @@ def _tiny_models(particle_type_cfg: dict | None = None):
         hidden_dim=16,
         n_res_blocks=1,
     )
-    target = (particle_type_cfg or {}).get("target", "physical")
+    resolved_type_cfg = particle_type_cfg or ParticleTypeConfig(target="physical")
+    target = resolved_type_cfg.target
     sec_dim = stage2_trunk_sec_dim(
-        particle_type_cfg or {"target": "physical"},
+        resolved_type_cfg,
         "flow",
         _K_MAX,
-        int(_PARTICLE_CFG["emb_dim"]),
+        _PARTICLE_CFG.emb_dim,
     )
     s2 = Stage2OneShot(
         pdg_vocab=3,
@@ -42,7 +44,7 @@ def _tiny_models(particle_type_cfg: dict | None = None):
         sec_dim=sec_dim,
         particle_type_cfg=particle_type_cfg,
     )
-    assert s2.particle_type_cfg.get("target", "physical") == target
+    assert s2.particle_type_cfg.target == target
     return s1.eval(), s2.eval()
 
 
@@ -95,7 +97,7 @@ def test_validate_marginals_physical_target_shapes():
 
 
 def test_validate_marginals_onehot_type_class_marginal():
-    particle_type_cfg = {"target": "onehot"}
+    particle_type_cfg = ParticleTypeConfig(target="onehot")
     s1, s2 = _tiny_models(particle_type_cfg)
     loader = _loader(n_sec_value=2, n_classes=s2.type_dim)
 

tests/test_wgan.py

@@ -1,12 +1,13 @@
 import torch
 
+from giant.config import ConditioningAxisConfig
 from giant.constants import COND_DIM, K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM
 from giant.model.network import CriticModel, Stage1Model, Stage2OneShot
 from giant.model.wgan import critic_loss, generator_loss, gradient_penalty
 from giant.sample import sample_secondaries_wgan, sample_wgan
 
-PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
-MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
+PARTICLE_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
+MATERIAL_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
 
 
 def _cond(B=8):

uv.lock

@@ -633,7 +633,7 @@ wheels = [
 
 [[package]]
 name = "giant"
-version = "0.3.0"
+version = "0.3.2"
 source = { editable = "." }
 dependencies = [
     { name = "numpy" },