v0.3.0 post-implementation audit: resolve all 9 tracked discrepancies

Works through docs/v0.3.0-followups.md item by item, closing the gap
between the design doc and the shipped v0.3.0-stage2-autoregressive code:

1. validate.py: 7-tuple batch unpacking, sample_stage1/sample_stage2
   dispatch, stage-2 particle-type-class marginal.
2. Stage-prefixed --stage1-*/--stage2-* CLI flags for train/new-run.
3. Thread stage2_model.k_max through loader/transforms/dataset/pipeline/
   train instead of the hardcoded K_MAX constant.
4. Mixed conditioning.particle.type / conditioning.material.type support
   end-to-end (data pipeline + dwarf warm-cache).
5. conditioning.share_stages = true: one shared ConditionEncoder instance
   across both stages.
6. stage2_model.generator = "ddpm" formally deferred into design doc §11.2
   (was silently unimplemented).
7. giant predict/rollout: implement conditioning.*.type = "onehot" via the
   checkpoint's saved pdg_topn_map/mat_topn_map.
8. network.py's checkpoint-path model_config migration now fails loudly on
   non-zero legacy expert_hidden_dim/expert_n_blocks, matching config.py's
   TOML-load path (§4.2).
9. validate_config now rejects stage2_model.n_sec.mode = "truth" for a
   rollout-capable checkpoint (§9).

Also cleared all pre-existing `ty check` noise (44 -> 0 diagnostics),
mostly a test-helper dict-unpack pattern that made every unrelated
constructor keyword look like a type error.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

docs/v0.3.0-design.md

@@ -892,6 +892,13 @@ All in `giant/config.py`:
   undesigned:** the mechanism should be worked out on its own terms when it is
   taken up, not pre-shaped by choices made for this refactor. Adding the block
   later is a config addition, not a break.
+- **`stage2_model.generator = "ddpm"`.** The value is **accepted by the
+  schema** (§3.3 lists `"flow" | "ddpm" | "wgan"` with no caveat) but
+  `FlowDDPMStageTrainer.__init__` (`giant/train.py`) raises
+  `NotImplementedError` for stage 2 — only `"flow"` and `"wgan"` have a
+  stage-2 secondary-decoder loss implemented. `stage1_model.generator =
+  "ddpm"` is unaffected; this restriction is stage-2-only. Landing stage-2
+  ddpm later is a trainer addition, not a config break.
 
 ### 11.3 Tracked as implementation work
 

docs/v0.3.0-followups.md

@@ -0,0 +1,37 @@
+# v0.3.0 — post-implementation audit: open discrepancies
+
+**Status:** steps 1–7 of `docs/v0.3.0-design.md` §12 are implemented (branch
+`v0.3.0-stage2-autoregressive`, commits `eb6dd27`..`200c6d2`). This document
+tracked discrepancies found between that implementation and the design contract
+during a 2026-08-07 audit, as concrete work items. **All items (1-9) are now
+resolved** — either implemented (1-5, 7-9) or explicitly deferred into
+`docs/v0.3.0-design.md` §11.2 (6: `stage2_model.generator = "ddpm"`). Step 8
+(`estimate_batch_size` recalibration) was intentionally still outstanding per
+§12 and was never tracked here.
+
+---
+
+## Confirmed correct during the audit (no action needed)
+
+For reference — these were explicitly checked against the design doc and
+match it, including two spots the doc itself flagged as likely stale that
+turned out fine:
+
+- Config schema, `DEFAULT_CONFIG`, `migrate_config` table (§4), `save_config`/
+  `merge_cli_overrides` recursion, `default_out_dir_name`, `Conditioning` enum,
+  `n_sec.mode = "stop_token"` error (§9, §11.2).
+- `network.py`'s full class decomposition (§5.3), dict-returning
+  `build_models`/`build_critics` (§5.4), `ExpertTrunk` separate in/out dims,
+  ST-Gumbel wiring (§2.1), AR token layout (§6.1), Markov/Attention history
+  encoders (§6.2).
+- Shared PDG top-N type map (one map, not two, per §8), `other_policy`
+  sample/modal/drop (§11.1), embedding L1-nearest decode, and the L1-distance
+  diagnostic surfaced in `giant analyze` (§11.3).
+- `analysis/render.py`/`analysis/router_gating.py` correctly branch
+  old-flat vs new-nested `model_config["router"]` location — doc flagged this
+  as a likely stale spot (§10) but it's actually fine.
+- `train.py`'s per-stage trainers, mixed flow+wgan runs, WGAN critic cadence,
+  per-stage router auxiliary losses, stage-prefixed metrics, stage-2-only
+  training via ground-truth `x1_s1` (§7).
+- `pipeline.py`'s deleted wgan+router rejection, per-stage `centers_init`
+  seeding, removed stale expert-size warning (§9, §10).

giant/analysis/router_gating.py

@@ -61,21 +61,26 @@ class _RouterHandle:
     pdg_map: dict[int, int]
     mat_map: dict[str, int]
     cond_normalizer: "Normalizer"
-    conditioning: str
+    particle_conditioning: str
+    material_conditioning: str
     router_type: str
 
 
-def _conditioning_str(model_cfg: dict, default: str = "embedding") -> str:
-    """The single conditioning-mode string `giant.data.transforms.
-    build_cond_features` still expects — from either a v0.2 checkpoint's
-    flat `model_config["conditioning"]` (already a string) or a new-format
-    one (`model_config["conditioning"]["particle"]["type"]`; `validate_config`
-    guarantees particle/material agree, see its mixed-conditioning check).
-    Mirrors `giant.cli._conditioning_str`."""
+def _conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]:
+    """(particle_conditioning, material_conditioning) for
+    `giant.data.transforms.build_cond_features` — from either a v0.2
+    checkpoint's flat `model_config["conditioning"]` (one shared string, same
+    for both axes) or a new-format one (independent
+    `model_config["conditioning"]["particle"/"material"]["type"]` —
+    docs/v0.3.0-design.md §3.1: the two axes may differ). Mirrors
+    `giant.cli._conditioning_axes`."""
     raw = model_cfg.get("conditioning", default)
     if isinstance(raw, dict):
-        return raw.get("particle", {}).get("type", default)
-    return raw
+        return (
+            raw.get("particle", {}).get("type", default),
+            raw.get("material", {}).get("type", default),
+        )
+    return raw, raw
 
 
 def load_router(checkpoint: str | Path) -> _RouterHandle | None:
@@ -107,12 +112,14 @@ def load_router(checkpoint: str | Path) -> _RouterHandle | None:
     if router is None:
         return None
 
+    particle_conditioning, material_conditioning = _conditioning_axes(model_cfg)
     return _RouterHandle(
         router=router,
         pdg_map={int(k): v for k, v in ckpt["pdg_map"].items()},
         mat_map={str(k): v for k, v in ckpt["mat_map"].items()},
         cond_normalizer=Normalizer.from_dict(ckpt["normalizer"]["cond"]),
-        conditioning=_conditioning_str(model_cfg),
+        particle_conditioning=particle_conditioning,
+        material_conditioning=material_conditioning,
         router_type=router_cfg["type"],
     )
 
@@ -169,7 +176,8 @@ def _gate_for_df(
         handle.pdg_map,
         handle.mat_map,
         cond_normalizer=handle.cond_normalizer,
-        conditioning=handle.conditioning,
+        particle_conditioning=handle.particle_conditioning,
+        material_conditioning=handle.material_conditioning,
     )
     with torch.no_grad():
         gate = handle.router.gate(

giant/cli.py

@@ -19,6 +19,7 @@ from tqdm import tqdm
 
 from giant import config as gconfig
 from giant.constants import (
+    K_MAX,
     LOCAL_TARGET_NAMES,
     PREDICT_COORD_METADATA_KEY,
     PREDICT_SCHEMA_VERSION,
@@ -72,44 +73,21 @@ def _router_total_experts(router_cfg: dict) -> int:
     return int(router_cfg.get("n_experts", 1))
 
 
-def _conditioning_str(model_cfg: dict, default: str = "embedding") -> str:
-    """The single conditioning-mode string `giant.data.transforms.
-    build_cond_features` still expects — from either a v0.2 checkpoint's
-    flat `model_config["conditioning"]` (already a string) or a new-format
-    one (`model_config["conditioning"]["particle"]["type"]`; `validate_config`
-    guarantees particle/material agree, see its mixed-conditioning check)."""
+def _conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]:
+    """(particle_conditioning, material_conditioning) for
+    `giant.data.transforms.build_cond_features`/`build_features` — from
+    either a v0.2 checkpoint's flat `model_config["conditioning"]` (one
+    shared string, same for both axes) or a new-format one (independent
+    `model_config["conditioning"]["particle"/"material"]["type"]` —
+    docs/v0.3.0-design.md §3.1: the two axes are configured independently and
+    may differ)."""
     raw = model_cfg.get("conditioning", default)
     if isinstance(raw, dict):
-        return raw.get("particle", {}).get("type", default)
-    return raw
-
-
-def _check_conditioning_onehot_support(model_cfg: dict, command: str) -> None:
-    """`giant predict`/`giant rollout` don't yet support
-    `conditioning.particle/material.type = "onehot"` — that needs
-    `pdg_topn_map.class_map`/`mat_topn_map.class_map` threaded into every
-    `build_cond_features` call in this module and `giant/rollout.py`'s
-    `_step_chunk`, which hasn't been wired (a separate axis from
-    `stage2_model.particle_type.target`, handled below — see
-    docs/v0.3.0-design.md §3.1 vs §3.3). Without this guard, predict/rollout
-    would silently build a `cond_cat` missing the topN columns
-    `ConditionEncoder`'s `"onehot"` mode expects. A v0.2 (flat) `model_config`
-    never has this, so this is a no-op there.
-    """
-    conditioning = model_cfg.get("conditioning")
-    if not isinstance(conditioning, dict):
-        return
-    particle_type = conditioning.get("particle", {}).get("type")
-    material_type = conditioning.get("material", {}).get("type")
-    if particle_type == "onehot" or material_type == "onehot":
-        typer.echo(
-            f"error: giant {command} does not yet support "
-            "conditioning.particle/material.type = 'onehot' — only "
-            "stage2_model.particle_type.target in ('onehot', 'embedding') "
-            "is implemented (docs/v0.3.0-design.md step 6).",
-            err=True,
+        return (
+            raw.get("particle", {}).get("type", default),
+            raw.get("material", {}).get("type", default),
         )
-        raise typer.Exit(1)
+    return raw, raw
 
 
 def _stage_cfg(model_cfg: dict, stage: str) -> dict:
@@ -144,6 +122,14 @@ def _load_pdg_topn_map(ckpt: dict):
     return topnmap_from_json(raw, axis="pdg") if raw is not None else None
 
 
+def _load_mat_topn_map(ckpt: dict):
+    """`ckpt["mat_topn_map"]` as a `giant.data.loader.TopNMap`, or `None` if
+    this checkpoint's `conditioning.material.type` was never `"onehot"` (see
+    `giant.pipeline.run_setup_stage`)."""
+    raw = ckpt.get("mat_topn_map")
+    return topnmap_from_json(raw, axis="material") if raw is not None else None
+
+
 def _batch_size_estimate_dims(
     model_cfg: dict, training: bool, stage: str = "stage1"
 ) -> tuple[int, int]:
@@ -295,6 +281,16 @@ class Mode(str, Enum):
     wgan = "wgan"
 
 
+class Decoder(str, Enum):
+    one_shot = "one_shot"
+    autoregressive = "autoregressive"
+
+
+class Stage1Context(str, Enum):
+    truth = "truth"
+    sampled = "sampled"
+
+
 # Conditioning itself lives in giant.config (imported below as gconfig) —
 # shared with scripts/dwarf.py's Typer commands so the two CLIs can't
 # silently drift apart on the option's valid values.
@@ -389,6 +385,89 @@ def train(
             "--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)"
         ),
     ] = None,
+    stage1_generator: Annotated[
+        Optional[Mode],
+        typer.Option(
+            "--stage1-generator",
+            help="Stage 1's generative objective — overrides --mode for stage 1 "
+            "only (see docs/v0.3.0-design.md decision 7)",
+        ),
+    ] = None,
+    stage1_hidden_dim: Annotated[
+        Optional[int],
+        typer.Option(
+            "--stage1-hidden-dim",
+            help="Overrides --hidden-dim for stage 1 only (same effect today; "
+            "--hidden-dim is kept as a shorthand since stage 1 was the only "
+            "target before stage2_model got its own flags)",
+        ),
+    ] = None,
+    stage1_n_res_blocks: Annotated[
+        Optional[int],
+        typer.Option(
+            "--stage1-n-res-blocks", help="Overrides --n-blocks for stage 1 only"
+        ),
+    ] = None,
+    stage1_dropout: Annotated[
+        Optional[float],
+        typer.Option("--stage1-dropout", help="Overrides --dropout for stage 1 only"),
+    ] = None,
+    stage2_generator: Annotated[
+        Optional[Mode],
+        typer.Option(
+            "--stage2-generator",
+            help="Stage 2's generative objective — overrides --mode for stage 2 "
+            "only, e.g. combine with --stage1-generator flow for a mixed "
+            "flow/wgan run",
+        ),
+    ] = None,
+    stage2_hidden_dim: Annotated[
+        Optional[int],
+        typer.Option("--stage2-hidden-dim", help="Stage 2 trunk width"),
+    ] = None,
+    stage2_n_res_blocks: Annotated[
+        Optional[int],
+        typer.Option("--stage2-n-res-blocks", help="Stage 2 trunk depth"),
+    ] = None,
+    stage2_dropout: Annotated[
+        Optional[float],
+        typer.Option("--stage2-dropout", help="Dropout inside stage 2's ResBlocks"),
+    ] = None,
+    stage2_decoder: Annotated[
+        Optional[Decoder],
+        typer.Option(
+            "--stage2-decoder",
+            help="one_shot: predict all k_max secondary slots at once (v0.2 "
+            "behaviour). autoregressive: emit one secondary at a time in "
+            "descending-energy order (default)",
+        ),
+    ] = None,
+    stage2_k_max: Annotated[
+        Optional[int],
+        typer.Option(
+            "--stage2-k-max",
+            help="Maximum secondary slots (fixed width under one_shot, a "
+            "generation-loop safety cap under autoregressive; default: 15)",
+        ),
+    ] = None,
+    stage2_context_dim: Annotated[
+        Optional[int],
+        typer.Option(
+            "--stage2-context-dim",
+            help="Width of the projected stage-1 outcome fed into stage 2's "
+            "conditioning (default: 64)",
+        ),
+    ] = None,
+    stage2_stage1_context: Annotated[
+        Optional[Stage1Context],
+        typer.Option(
+            "--stage2-stage1-context",
+            help="What stage 2 conditions on during training: 'truth' (the "
+            "ground-truth stage-1 target, detached — default) or 'sampled' "
+            "(stage 1's own sampled output, closing the train/inference gap "
+            "at the cost of an extra sampling pass per batch)",
+        ),
+    ] = None,
     conditioning: Annotated[
         Optional[Conditioning],
         typer.Option(
@@ -457,6 +536,52 @@ def train(
             "(default: same as --lr)",
         ),
     ] = None,
+    stage1_n_critic: Annotated[
+        Optional[int],
+        typer.Option("--stage1-n-critic", help="Overrides --n-critic for stage 1 only"),
+    ] = None,
+    stage1_gp_weight: Annotated[
+        Optional[float],
+        typer.Option(
+            "--stage1-gp-weight", help="Overrides --gp-weight for stage 1 only"
+        ),
+    ] = None,
+    stage1_noise_dim: Annotated[
+        Optional[int],
+        typer.Option(
+            "--stage1-noise-dim", help="Overrides --noise-dim for stage 1 only"
+        ),
+    ] = None,
+    stage1_critic_lr: Annotated[
+        Optional[float],
+        typer.Option(
+            "--stage1-critic-lr", help="Overrides --critic-lr for stage 1 only"
+        ),
+    ] = None,
+    stage2_n_critic: Annotated[
+        Optional[int],
+        typer.Option("--stage2-n-critic", help="Overrides --n-critic for stage 2 only"),
+    ] = None,
+    stage2_gp_weight: Annotated[
+        Optional[float],
+        typer.Option(
+            "--stage2-gp-weight", help="Overrides --gp-weight for stage 2 only"
+        ),
+    ] = None,
+    stage2_noise_dim: Annotated[
+        Optional[int],
+        typer.Option(
+            "--stage2-noise-dim",
+            help="Overrides --noise-dim for stage 2 only; under "
+            "--stage2-decoder autoregressive a fresh draw is made per token",
+        ),
+    ] = None,
+    stage2_critic_lr: Annotated[
+        Optional[float],
+        typer.Option(
+            "--stage2-critic-lr", help="Overrides --critic-lr for stage 2 only"
+        ),
+    ] = None,
     val_fraction: Annotated[
         Optional[float], typer.Option("--val-fraction", "-f")
     ] = None,
@@ -597,9 +722,11 @@ def train(
         }.items()
         if v is not None
     }
-    # Architecture-shape flags apply to stage1_model only (decision: stage 2
-    # is tuned via a hand-edited config.toml until stage-prefixed --stage2-*
-    # flags land — see docs/v0.3.0-design.md decision 7).
+    # --hidden-dim/--n-blocks/--dropout are stage-1-only shorthands kept for
+    # backward compatibility (they predate stage2_model having its own
+    # flags); --stage1-*/--stage2-* below are the explicit, discoverable
+    # per-stage flags docs/v0.3.0-design.md decision 7 calls for, and take
+    # precedence when both are given.
     cli_stage1_model: dict[str, object] = {
         k: v
         for k, v in {
@@ -609,6 +736,32 @@ def train(
         }.items()
         if v is not None
     }
+    cli_stage1_model.update(
+        {
+            k: v
+            for k, v in {
+                "hidden_dim": stage1_hidden_dim,
+                "n_res_blocks": stage1_n_res_blocks,
+                "dropout": stage1_dropout,
+            }.items()
+            if v is not None
+        }
+    )
+    cli_stage2_model: dict[str, object] = {
+        k: v
+        for k, v in {
+            "hidden_dim": stage2_hidden_dim,
+            "n_res_blocks": stage2_n_res_blocks,
+            "dropout": stage2_dropout,
+            "decoder": stage2_decoder.value if stage2_decoder is not None else None,
+            "k_max": stage2_k_max,
+            "context_dim": stage2_context_dim,
+            "stage1_context": stage2_stage1_context.value
+            if stage2_stage1_context is not None
+            else None,
+        }.items()
+        if v is not None
+    }
     cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
     if cli_router:
         cli_stage1_model["router"] = cli_router
@@ -628,37 +781,65 @@ def train(
         overrides["train"] = cli_train
     if cli_stage1_model:
         overrides["stage1_model"] = cli_stage1_model
+    if cli_stage2_model:
+        overrides["stage2_model"] = cli_stage2_model
     if cli_conditioning:
         overrides["conditioning"] = cli_conditioning
 
-    # --mode/--n-critic/--gp-weight/--critic-lr apply to BOTH stages — v0.2
-    # had one global mode/wgan config shared by both (see
-    # giant.config.migrate_config's train.mode /
-    # train.{n_critic,gp_weight,critic_lr} precedent), unlike the
-    # architecture-shape flags above.
+    # --mode/--n-critic/--gp-weight/--critic-lr/--noise-dim apply to BOTH
+    # stages by default (v0.2 had one global mode/wgan config shared by both
+    # — see giant.config.migrate_config's train.mode /
+    # train.{n_critic,gp_weight,critic_lr} precedent); the --stage1-*/
+    # --stage2-* variants below override a single stage independently (decision
+    # 7), which is what actually enables e.g. `--stage1-generator flow
+    # --stage2-generator wgan`.
     if mode is not None:
         overrides.setdefault("stage1_model", {})["generator"] = mode.value
         overrides.setdefault("stage2_model", {})["generator"] = mode.value
-    if noise_dim is not None:
-        overrides.setdefault("stage1_model", {}).setdefault("wgan", {})["noise_dim"] = (
-            noise_dim
-        )
-    wgan_overrides = {
+    if stage1_generator is not None:
+        overrides.setdefault("stage1_model", {})["generator"] = stage1_generator.value
+    if stage2_generator is not None:
+        overrides.setdefault("stage2_model", {})["generator"] = stage2_generator.value
+
+    shared_wgan_overrides = {
         k: v
         for k, v in {
             "n_critic": n_critic,
             "gp_weight": gp_weight,
+            "noise_dim": noise_dim,
             "critic_lr": critic_lr,
         }.items()
         if v is not None
     }
-    if wgan_overrides:
-        overrides.setdefault("stage1_model", {}).setdefault("wgan", {}).update(
-            wgan_overrides
-        )
-        overrides.setdefault("stage2_model", {}).setdefault("wgan", {}).update(
-            wgan_overrides
-        )
+    stage1_wgan_overrides = {
+        k: v
+        for k, v in {
+            "n_critic": stage1_n_critic,
+            "gp_weight": stage1_gp_weight,
+            "noise_dim": stage1_noise_dim,
+            "critic_lr": stage1_critic_lr,
+        }.items()
+        if v is not None
+    }
+    stage2_wgan_overrides = {
+        k: v
+        for k, v in {
+            "n_critic": stage2_n_critic,
+            "gp_weight": stage2_gp_weight,
+            "noise_dim": stage2_noise_dim,
+            "critic_lr": stage2_critic_lr,
+        }.items()
+        if v is not None
+    }
+    for stage_name, stage_specific in (
+        ("stage1_model", stage1_wgan_overrides),
+        ("stage2_model", stage2_wgan_overrides),
+    ):
+        stage_wgan = {**shared_wgan_overrides, **stage_specific}
+        if stage_wgan:
+            overrides.setdefault(stage_name, {}).setdefault("wgan", {}).update(
+                stage_wgan
+            )
 
     cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides)
     gconfig.validate_config(cfg)
@@ -735,6 +916,36 @@ def new_run(
     n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None,
     emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None,
     dropout: Annotated[Optional[float], typer.Option("--dropout", "-d")] = None,
+    stage1_generator: Annotated[
+        Optional[Mode], typer.Option("--stage1-generator")
+    ] = None,
+    stage1_hidden_dim: Annotated[
+        Optional[int], typer.Option("--stage1-hidden-dim")
+    ] = None,
+    stage1_n_res_blocks: Annotated[
+        Optional[int], typer.Option("--stage1-n-res-blocks")
+    ] = None,
+    stage1_dropout: Annotated[Optional[float], typer.Option("--stage1-dropout")] = None,
+    stage2_generator: Annotated[
+        Optional[Mode], typer.Option("--stage2-generator")
+    ] = None,
+    stage2_hidden_dim: Annotated[
+        Optional[int], typer.Option("--stage2-hidden-dim")
+    ] = None,
+    stage2_n_res_blocks: Annotated[
+        Optional[int], typer.Option("--stage2-n-res-blocks")
+    ] = None,
+    stage2_dropout: Annotated[Optional[float], typer.Option("--stage2-dropout")] = None,
+    stage2_decoder: Annotated[
+        Optional[Decoder], typer.Option("--stage2-decoder")
+    ] = None,
+    stage2_k_max: Annotated[Optional[int], typer.Option("--stage2-k-max")] = None,
+    stage2_context_dim: Annotated[
+        Optional[int], typer.Option("--stage2-context-dim")
+    ] = None,
+    stage2_stage1_context: Annotated[
+        Optional[Stage1Context], typer.Option("--stage2-stage1-context")
+    ] = None,
     conditioning: Annotated[
         Optional[Conditioning], typer.Option("--conditioning")
     ] = None,
@@ -802,6 +1013,32 @@ def new_run(
         }.items()
         if v is not None
     }
+    cli_stage1_model.update(
+        {
+            k: v
+            for k, v in {
+                "hidden_dim": stage1_hidden_dim,
+                "n_res_blocks": stage1_n_res_blocks,
+                "dropout": stage1_dropout,
+            }.items()
+            if v is not None
+        }
+    )
+    cli_stage2_model: dict[str, object] = {
+        k: v
+        for k, v in {
+            "hidden_dim": stage2_hidden_dim,
+            "n_res_blocks": stage2_n_res_blocks,
+            "dropout": stage2_dropout,
+            "decoder": stage2_decoder.value if stage2_decoder is not None else None,
+            "k_max": stage2_k_max,
+            "context_dim": stage2_context_dim,
+            "stage1_context": stage2_stage1_context.value
+            if stage2_stage1_context is not None
+            else None,
+        }.items()
+        if v is not None
+    }
     cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
     if cli_router:
         cli_stage1_model["router"] = cli_router
@@ -818,11 +1055,17 @@ def new_run(
         overrides["train"] = cli_train
     if cli_stage1_model:
         overrides["stage1_model"] = cli_stage1_model
+    if cli_stage2_model:
+        overrides["stage2_model"] = cli_stage2_model
     if cli_conditioning:
         overrides["conditioning"] = cli_conditioning
     if mode is not None:
         overrides.setdefault("stage1_model", {})["generator"] = mode.value
         overrides.setdefault("stage2_model", {})["generator"] = mode.value
+    if stage1_generator is not None:
+        overrides.setdefault("stage1_model", {})["generator"] = stage1_generator.value
+    if stage2_generator is not None:
+        overrides.setdefault("stage2_model", {})["generator"] = stage2_generator.value
 
     cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides)
     gconfig.validate_config(cfg)
@@ -988,10 +1231,11 @@ def predict(
         raise typer.Exit(1)
 
     model_cfg = ckpt["model_config"]
-    _check_conditioning_onehot_support(model_cfg, "predict")
     pdg_topn_map = _load_pdg_topn_map(ckpt)
+    mat_topn_map = _load_mat_topn_map(ckpt)
     other_policy = _particle_type_other_policy(model_cfg)
     stage1_ddpm_steps = _ddpm_steps(model_cfg, "stage1")
+    stage2_k_max = _stage_cfg(model_cfg, "stage2").get("k_max", K_MAX)
 
     if batch_size_auto:
         est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(
@@ -1013,7 +1257,21 @@ def predict(
 
     assert batch_size_value is not None
     bs = batch_size_value
-    conditioning = _conditioning_str(model_cfg)
+    particle_conditioning, material_conditioning = _conditioning_axes(model_cfg)
+    if particle_conditioning == "onehot" and pdg_topn_map is None:
+        typer.echo(
+            "error: checkpoint's conditioning.particle.type='onehot' but has "
+            "no pdg_topn_map — retrain with the current code",
+            err=True,
+        )
+        raise typer.Exit(1)
+    if material_conditioning == "onehot" and mat_topn_map is None:
+        typer.echo(
+            "error: checkpoint's conditioning.material.type='onehot' but has "
+            "no mat_topn_map — retrain with the current code",
+            err=True,
+        )
+        raise typer.Exit(1)
     pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
     mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
     cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
@@ -1050,7 +1308,18 @@ def predict(
     skipped = 0
     unknown_pdg_counts: Counter[int] = Counter()
     total_rows = sum(pq.ParquetFile(path).metadata.num_rows for path in files)
-    chunk_iter = iter_file_chunks if coord == Coord.local else iter_cond_chunks
+
+    def chunk_iter(path: Path, offset: int):
+        if coord == Coord.local:
+            return iter_file_chunks(path, offset=offset, k_max=stage2_k_max)
+        return iter_cond_chunks(path, offset=offset)
+
+    cond_pdg_topn = (
+        pdg_topn_map.class_map if particle_conditioning == "onehot" else None
+    )
+    cond_mat_topn = (
+        mat_topn_map.class_map if material_conditioning == "onehot" else None
+    )
 
     def _concat(
         a: dict[str, np.ndarray], b: dict[str, np.ndarray]
@@ -1062,12 +1331,26 @@ def predict(
 
         if coord == Coord.local:
             cond_cont, cond_cat, target_raw, _, _, _, _, _, _ = build_features(
-                piece, pdg_map, mat_map, conditioning=conditioning
+                piece,
+                pdg_map,
+                mat_map,
+                particle_conditioning=particle_conditioning,
+                material_conditioning=material_conditioning,
+                pdg_topn_map=cond_pdg_topn,
+                mat_topn_map=cond_mat_topn,
+                k_max=stage2_k_max,
             )
             cond_cont = cond_norm.transform(cond_cont)
         else:
             cond_cont, cond_cat = build_cond_features(
-                piece, pdg_map, mat_map, cond_norm, conditioning=conditioning
+                piece,
+                pdg_map,
+                mat_map,
+                cond_norm,
+                particle_conditioning=particle_conditioning,
+                material_conditioning=material_conditioning,
+                pdg_topn_map=cond_pdg_topn,
+                mat_topn_map=cond_mat_topn,
             )
 
         cc = torch.from_numpy(cond_cont).float().to(_device)
@@ -1409,9 +1692,23 @@ def rollout(
     training_cfg = gconfig.load_checkpoint_config(checkpoint)
 
     model_cfg = ckpt["model_config"]
-    _check_conditioning_onehot_support(model_cfg, "rollout")
-    conditioning = _conditioning_str(model_cfg)
+    particle_conditioning, material_conditioning = _conditioning_axes(model_cfg)
     pdg_topn_map = _load_pdg_topn_map(ckpt)
+    mat_topn_map = _load_mat_topn_map(ckpt)
+    if particle_conditioning == "onehot" and pdg_topn_map is None:
+        typer.echo(
+            "error: checkpoint's conditioning.particle.type='onehot' but has "
+            "no pdg_topn_map — retrain with the current code",
+            err=True,
+        )
+        raise typer.Exit(1)
+    if material_conditioning == "onehot" and mat_topn_map is None:
+        typer.echo(
+            "error: checkpoint's conditioning.material.type='onehot' but has "
+            "no mat_topn_map — retrain with the current code",
+            err=True,
+        )
+        raise typer.Exit(1)
     other_policy = _particle_type_other_policy(model_cfg)
     stage1_ddpm_steps = _ddpm_steps(model_cfg, "stage1")
     stage2_ddpm_steps = _ddpm_steps(model_cfg, "stage2")
@@ -1489,8 +1786,10 @@ def rollout(
         max_tracks_per_event=max_tracks_per_event,
         escape_threshold=escape_threshold,
         on_chunk=_write_chunk,
-        conditioning=conditioning,
+        particle_conditioning=particle_conditioning,
+        material_conditioning=material_conditioning,
         pdg_topn_map=pdg_topn_map,
+        mat_topn_map=mat_topn_map,
         other_policy=other_policy,
         seed=seed,
         stage1_ddpm_steps=stage1_ddpm_steps,

giant/config.py

@@ -680,18 +680,6 @@ def validate_config(cfg: dict) -> None:
     single block's defaults.
     """
     particle_type = _get_path(cfg, "conditioning.particle.type")
-    material_type = _get_path(cfg, "conditioning.material.type")
-
-    if particle_type != material_type:
-        raise ValueError(
-            "conditioning.particle.type "
-            f"({particle_type!r}) != conditioning.material.type "
-            f"({material_type!r}) — the data pipeline (giant/data/transforms.py's "
-            "build_features/build_cond_features) doesn't support mixed "
-            "particle/material conditioning types yet, even though "
-            "ConditionEncoder itself (giant/model/network.py) is already able "
-            "to. Not implemented in v0.3.0 — use the same type for both axes."
-        )
 
     pt_target = _get_path(cfg, "stage2_model.particle_type.target")
     if pt_target == "embedding" and particle_type != "embedding":
@@ -733,6 +721,21 @@ def validate_config(cfg: dict) -> None:
             "(standalone stage-2 evaluation only, never for rollout)"
         )
 
+    if (
+        _get_path(cfg, "stage2_model.n_sec.mode") == "truth"
+        and _get_path(cfg, "stage1_model.active")
+        and _get_path(cfg, "stage2_model.active")
+    ):
+        raise ValueError(
+            "stage2_model.n_sec.mode = 'truth' is invalid for a "
+            "rollout-capable checkpoint (both stage1_model.active and "
+            "stage2_model.active = true — see docs/v0.3.0-design.md §9): "
+            "'truth' takes n_sec from ground truth, which giant rollout "
+            "doesn't have. 'truth' is for standalone stage-2 evaluation "
+            "only — set stage1_model.active = false for that, or use "
+            "'head' (default) for a rollout-capable checkpoint."
+        )
+
     if _get_path(cfg, "stage2_model.decoder") == "autoregressive":
         history = _get_path(cfg, "stage2_model.autoregressive.history")
         if history not in ("markov", "attention"):

giant/data/dataset.py

@@ -6,6 +6,7 @@ import numpy as np
 import torch
 from torch.utils.data import IterableDataset
 
+from giant.constants import K_MAX
 from giant.data.loader import event_id_offset, iter_file_chunks
 from giant.data.transforms import Normalizer, build_features, sorted_membership
 
@@ -46,7 +47,7 @@ class StreamingStepsDataset(IterableDataset):
                     (see docs/v0.3.0-design.md decision 4)
         target_s1:  (B, 9) float32 — normalised Stage-1 primary target
         n_sec:      (B,) int64 — true secondary count per step
-        sec_cont:   (B, K_MAX, SEC_SLOT_DIM) float32 — [stick_logit,
+        sec_cont:   (B, k_max, SEC_SLOT_DIM) float32 — [stick_logit,
                     local_dir, log_mass, charge] per slot (mass/charge
                     normalised iff `sec_phys_normalizer` was given); always
                     computed the same way regardless of
@@ -54,9 +55,13 @@ class StreamingStepsDataset(IterableDataset):
                     downstream under target="physical"
         proc_idx:   (B,) int64 — process-class label (ProcessRouter supervision
                     only; zeros when `proc_map` is None)
-        sec_type_idx: (B, K_MAX) int64 — per-slot class index into
+        sec_type_idx: (B, k_max) int64 — per-slot class index into
                     `sec_type_class_map`, for particle_type.target in
                     ("onehot", "embedding"); zeros (unused) otherwise
+
+    `k_max` (constructor arg, default the module constant) should match
+    `stage2_model.k_max` (docs/v0.3.0-design.md §8) — it sets the padded
+    width of `sec_cont`/`sec_type_idx` above.
     """
 
     def __init__(
@@ -71,11 +76,13 @@ class StreamingStepsDataset(IterableDataset):
         shuffle_buffer: int = 65536,
         shuffle: bool = True,
         proc_map: dict[str, int] | None = None,
-        conditioning: str = "embedding",
+        particle_conditioning: str = "embedding",
+        material_conditioning: str = "embedding",
         sec_phys_normalizer: Normalizer | None = None,
         pdg_topn_map: dict[int, int] | None = None,
         mat_topn_map: dict[str, int] | None = None,
         sec_type_class_map: dict | None = None,
+        k_max: int = K_MAX,
     ) -> None:
         self.files = list(files)
         self._offsets = {path: event_id_offset(i) for i, path in enumerate(self.files)}
@@ -89,11 +96,13 @@ class StreamingStepsDataset(IterableDataset):
         self.shuffle_buffer = max(shuffle_buffer, batch_size)
         self.shuffle = shuffle
         self.proc_map = proc_map
-        self.conditioning = conditioning
+        self.particle_conditioning = particle_conditioning
+        self.material_conditioning = material_conditioning
         self.sec_phys_normalizer = sec_phys_normalizer
         self.pdg_topn_map = pdg_topn_map
         self.mat_topn_map = mat_topn_map
         self.sec_type_class_map = sec_type_class_map
+        self.k_max = k_max
 
     def __iter__(self):
         worker_info = torch.utils.data.get_worker_info()
@@ -115,7 +124,9 @@ class StreamingStepsDataset(IterableDataset):
         buf_n = 0
 
         for path in files:
-            for chunk in iter_file_chunks(path, offset=self._offsets[path]):
+            for chunk in iter_file_chunks(
+                path, offset=self._offsets[path], k_max=self.k_max
+            ):
                 mask = sorted_membership(chunk["event_id"], self._events_arr)
                 if not mask.any():
                     continue
@@ -140,10 +151,12 @@ class StreamingStepsDataset(IterableDataset):
                     sec_phys_normalizer=self.sec_phys_normalizer,
                     proc_map=self.proc_map,
                     require_secondaries=True,
-                    conditioning=self.conditioning,
+                    particle_conditioning=self.particle_conditioning,
+                    material_conditioning=self.material_conditioning,
                     pdg_topn_map=self.pdg_topn_map,
                     mat_topn_map=self.mat_topn_map,
                     sec_type_class_map=self.sec_type_class_map,
+                    k_max=self.k_max,
                 )
                 buf_cont.append(cond_cont)
                 buf_cat.append(cond_cat)

giant/data/loader.py

@@ -6,6 +6,8 @@ import numpy as np
 import pandas as pd
 import pyarrow.parquet as pq
 
+from giant.constants import K_MAX
+
 # A manifest is a plain text file listing one parquet path per line, used to
 # name a curated subset of files (e.g. a train/holdout pool) without copying
 # or symlinking the underlying parquet files. Lines are resolved relative to
@@ -113,9 +115,9 @@ def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndar
     return out
 
 
-def _df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
-    from giant.constants import K_MAX
-
+def _df_to_dict(
+    df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX
+) -> dict[str, np.ndarray]:
     has_sec_lists = "sec_E_list" in df.columns
 
     d: dict[str, np.ndarray] = {
@@ -147,17 +149,19 @@ def _df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
     }
 
     if has_sec_lists:
-        d["sec_E_list"] = _pad_list_col(df["sec_E_list"], K_MAX)
-        d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], K_MAX)
+        d["sec_E_list"] = _pad_list_col(df["sec_E_list"], k_max)
+        d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], k_max)
         d["sec_dir_list"] = _pad_dir_col(
-            df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], K_MAX
+            df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], k_max
         )
 
     return d
 
 
-def load_steps(path: str | Path, offset: int = 0) -> dict[str, np.ndarray]:
-    return _df_to_dict(pd.read_parquet(path), offset=offset)
+def load_steps(
+    path: str | Path, offset: int = 0, k_max: int = K_MAX
+) -> dict[str, np.ndarray]:
+    return _df_to_dict(pd.read_parquet(path), offset=offset, k_max=k_max)
 
 
 def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
@@ -167,12 +171,16 @@ def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
 
 
 def iter_file_chunks(
-    path: str | Path, offset: int = 0
+    path: str | Path, offset: int = 0, k_max: int = K_MAX
 ) -> Iterator[dict[str, np.ndarray]]:
-    """Yield one parquet row-group at a time so a large file never fully loads."""
+    """Yield one parquet row-group at a time so a large file never fully loads.
+
+    `k_max` sets the padded width of the sec_*_list columns (should match
+    `stage2_model.k_max` — see docs/v0.3.0-design.md §8); defaults to the
+    module constant for callers that don't care (e.g. Stage-1-only reads)."""
     pf = pq.ParquetFile(path)
     for i in range(pf.num_row_groups):
-        yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset)
+        yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset, k_max=k_max)
 
 
 _COND_COLS = [

giant/data/setup_cache.py

@@ -96,10 +96,22 @@ def fingerprint_files(files: list[Path]) -> list[list]:
     return out
 
 
-def normalizer_key(val_fraction: float, seed: int, conditioning: str) -> str:
+def normalizer_key(
+    val_fraction: float,
+    seed: int,
+    particle_conditioning: str,
+    material_conditioning: str,
+) -> str:
     # .6g avoids float-repr drift (e.g. 0.1 vs 0.10000000000000002) causing
-    # spurious cache misses between runs with the "same" val_fraction.
-    return f"valfrac={val_fraction:.6g}_seed={seed}_cond={conditioning}"
+    # spurious cache misses between runs with the "same" val_fraction. The two
+    # conditioning axes are independent (docs/v0.3.0-design.md §3.1) and both
+    # affect which cond_cont columns are computed for real vs. zero-filled
+    # (giant.data.transforms._physical_cond_columns), so both must be part of
+    # the key or two mixed-axis runs could collide on the same cache entry.
+    return (
+        f"valfrac={val_fraction:.6g}_seed={seed}_pcond={particle_conditioning}"
+        f"_mcond={material_conditioning}"
+    )
 
 
 # Top-N-map axes (docs/v0.3.0-design.md §8): "pdg" keys match pdg_map's int

giant/data/transforms.py

@@ -2,6 +2,8 @@ import warnings
 
 import numpy as np
 
+from giant.constants import K_MAX
+
 _EPS = 1e-8
 
 # Floor added to each energy fraction before taking log-ratios so the simplex
@@ -719,14 +721,23 @@ def decode_secondaries(
 
 
 def _physical_cond_columns(
-    data: dict[str, np.ndarray], conditioning: str
+    data: dict[str, np.ndarray],
+    particle_conditioning: str,
+    material_conditioning: str,
 ) -> np.ndarray:
     """(N, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM) physical conditioning columns.
 
-    "embedding"/"onehot" modes zero-fill (cheap, and ConditionEncoder never
-    reads these columns in either mode — so an unfilled giant.materials table
-    can never crash an "embedding"/"onehot"-mode run). "physical" mode
-    computes them for real: particle columns come from
+    The particle and material blocks are gated independently
+    (docs/v0.3.0-design.md §3.1: "configured independently and may mix
+    freely — e.g. material `physical` with particle `embedding`"), so e.g.
+    `particle_conditioning="embedding"` + `material_conditioning="physical"`
+    zero-fills only the particle columns and computes the material ones for
+    real.
+
+    "embedding"/"onehot" zero-fill their block (cheap, and ConditionEncoder
+    never reads these columns in either mode — so an unfilled
+    giant.materials table can never crash an "embedding"/"onehot"-mode run).
+    "physical" computes it for real: particle columns come from
     `data["mass"]`/`data["charge"]` when the caller already knows them
     directly (rollout.py, for a track descended from a model-predicted
     secondary — see giant/rollout.py's "no snapping" design), else derived
@@ -736,36 +747,47 @@ def _physical_cond_columns(
     """
     from giant.constants import MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM
 
-    if conditioning in ("embedding", "onehot"):
-        n = len(next(iter(data.values())))
-        return np.zeros((n, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM), dtype=np.float32)
-    if conditioning != "physical":
-        raise ValueError(f"unknown conditioning mode {conditioning!r}")
+    n = len(next(iter(data.values())))
 
-    from giant.materials import material_properties_array
-    from giant.particles import particle_phys_array
+    if particle_conditioning == "physical":
+        from giant.particles import particle_phys_array
 
-    if "mass" in data and "charge" in data:
-        mass = np.asarray(data["mass"], dtype=np.float32)
-        charge = np.asarray(data["charge"], dtype=np.float32)
+        if "mass" in data and "charge" in data:
+            mass = np.asarray(data["mass"], dtype=np.float32)
+            charge = np.asarray(data["charge"], dtype=np.float32)
+        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:
-        mass, charge = particle_phys_array(data["pdg"]).T
+        raise ValueError(
+            f"unknown conditioning.particle.type {particle_conditioning!r}"
+        )
 
-    z_eff, a_eff, density, x0, lambda_int = material_properties_array(
-        data["material"]
-    ).T
+    if material_conditioning == "physical":
+        from giant.materials import material_properties_array
 
-    return np.column_stack(
-        [
-            log_transform(mass),
-            charge,
-            z_eff,
-            a_eff,
-            log_transform(density),
-            log_transform(x0),
-            log_transform(lambda_int),
-        ]
-    ).astype(np.float32)
+        z_eff, a_eff, density, x0, lambda_int = material_properties_array(
+            data["material"]
+        ).T
+        material_cols = np.column_stack(
+            [
+                z_eff,
+                a_eff,
+                log_transform(density),
+                log_transform(x0),
+                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}"
+        )
+
+    return np.column_stack([particle_cols, material_cols]).astype(np.float32)
 
 
 def build_cond_features(
@@ -773,19 +795,24 @@ def build_cond_features(
     pdg_map: dict[int, int],
     mat_map: dict[str, int],
     cond_normalizer: "Normalizer | None" = None,
-    conditioning: str = "embedding",
+    particle_conditioning: str = "embedding",
+    material_conditioning: str = "embedding",
     pdg_topn_map: dict[int, int] | None = None,
     mat_topn_map: dict[str, int] | None = None,
 ) -> tuple[np.ndarray, np.ndarray]:
     """Build conditioning arrays only — no target, no post-step variables.
 
+    `particle_conditioning`/`material_conditioning` are independent
+    (docs/v0.3.0-design.md §3.1) — e.g. `particle_conditioning="embedding"` +
+    `material_conditioning="physical"` is a valid mix.
+
     `pdg_topn_map`/`mat_topn_map` (a top-N-plus-other `class_map`, see
     `giant.data.loader.build_topn_map_from_files`) append extra `cond_cat`
     columns read by `ConditionEncoder`'s `"onehot"` mode
     (docs/v0.3.0-design.md decision 4): pdg topN index at column 2 (iff
     `pdg_topn_map` given), material topN index at column 3 (iff
     `mat_topn_map` given, after column 2 if both are). Only ever given when
-    `conditioning == "onehot"`; `cond_cat` stays `(N, 2)` otherwise.
+    the corresponding axis is `"onehot"`; `cond_cat` stays `(N, 2)` otherwise.
     """
     cond_cont = np.column_stack(
         [
@@ -796,7 +823,10 @@ def build_cond_features(
         ]
     ).astype(np.float32)
     cond_cont = np.column_stack(
-        [cond_cont, _physical_cond_columns(data, conditioning)]
+        [
+            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
@@ -807,10 +837,11 @@ def build_cond_features(
     # 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".
-    strict = conditioning == "embedding"
-    pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=strict)
-    mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=strict)
+    # 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))
@@ -819,34 +850,44 @@ def build_cond_features(
     cond_cat = np.column_stack(cat_cols)
 
     if cond_normalizer is not None:
-        cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, conditioning)
+        cond_cont = _cond_normalizer_transform(
+            cond_cont, cond_normalizer, particle_conditioning, material_conditioning
+        )
 
     return cond_cont, cond_cat
 
 
 def _cond_normalizer_transform(
-    cond_cont: np.ndarray, cond_normalizer: "Normalizer", conditioning: str
+    cond_cont: np.ndarray,
+    cond_normalizer: "Normalizer",
+    particle_conditioning: str,
+    material_conditioning: str,
 ) -> np.ndarray:
     """Apply ``cond_normalizer``, padding a legacy narrower normalizer if needed.
 
     Checkpoints trained before physical-property conditioning (``COND_DIM``
     8->15, ``giant/constants.py``) saved a ``COND_DIM_BASE``-wide (8) cond
     normalizer, fit before ``build_cond_features`` grew the extra physical
-    columns. In "embedding" mode those columns are never read by
+    columns. When NEITHER axis is "physical" those columns are never read by
     ``ConditionEncoder`` (``giant/model/network.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. In
-    "physical" mode the physical columns are load-bearing, so a mismatch
-    there is a real incompatibility, not something to paper over.
+    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
+    real incompatibility, not something to paper over.
     """
     mean, std = cond_normalizer.mean, cond_normalizer.std
     assert mean is not None and std is not None, "Normalizer not fitted"
     width = cond_cont.shape[-1]
     if mean.shape[-1] < width:
-        if conditioning != "embedding":
+        physical_load_bearing = "physical" in (
+            particle_conditioning,
+            material_conditioning,
+        )
+        if physical_load_bearing:
             raise ValueError(
                 f"cond normalizer has {mean.shape[-1]} columns, expected "
-                f"{width}, and conditioning={conditioning!r} reads the "
+                f"{width}, and particle_conditioning={particle_conditioning!r}/"
+                f"material_conditioning={material_conditioning!r} reads the "
                 "physical columns directly — this checkpoint predates "
                 "physical-property conditioning and can't be safely padded; "
                 "retrain it under the current code."
@@ -867,11 +908,13 @@ def build_features(
     fit: bool = False,
     proc_map: dict[str, int] | None = None,
     require_secondaries: bool = False,
-    conditioning: str = "embedding",
+    particle_conditioning: str = "embedding",
+    material_conditioning: str = "embedding",
     sec_phys_only: bool = False,
     pdg_topn_map: dict[int, int] | None = None,
     mat_topn_map: dict[str, int] | None = None,
     sec_type_class_map: dict | None = None,
+    k_max: int = K_MAX,
 ) -> tuple[
     np.ndarray,
     np.ndarray,
@@ -922,8 +965,14 @@ def build_features(
         top-N-plus-other map's `class_map` for `target = "onehot"`, or the
         dense `pdg_map` for `target = "embedding"` (pass `pdg_map` itself).
         `None` for `target = "physical"`.
+
+    k_max: should match `stage2_model.k_max` (docs/v0.3.0-design.md §8) —
+        overridden internally by `data["sec_E_list"]`'s own padded width when
+        present (the loader already padded it to some k_max; that width is
+        authoritative), so this only actually matters when secondary list
+        columns are absent (Stage-1-only reads, or a pre-secondary-join
+        file), where it sets `sec_cont`/`sec_type_idx`'s zero-filled width.
     """
-    from giant.constants import K_MAX
 
     post_dir_local = local_frame_rotation(data["pre_dir"], data["post_dir"])
     travel_dir_local = local_frame_rotation(
@@ -953,7 +1002,10 @@ def build_features(
         ]
     ).astype(np.float32)  # (N, COND_DIM_BASE=8)
     cond_cont = np.column_stack(
-        [cond_cont, _physical_cond_columns(data, conditioning)]
+        [
+            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)
@@ -968,20 +1020,27 @@ def build_features(
     n_sec_raw = data["n_sec"].astype(
         np.int64
     )  # (N,) unclamped, for the valid-slot mask
-    # Clamp the classification label to K_MAX: the head only has K_MAX+1 classes
-    # (0..K_MAX), and truncating here mirrors the K_MAX-slot truncation already
-    # applied to sec_cont by the loader's list padding. Without this, a rare
-    # high-multiplicity step (real data goes up to ~37) hands cross_entropy
-    # an out-of-range target and CUDA asserts.
-    n_sec = np.minimum(n_sec_raw, K_MAX).astype(np.int64)  # (N,)
 
     # Secondary continuous targets
     sec_E_list = data.get("sec_E_list")
     sec_dir_list = data.get("sec_dir_list")
     sec_pdg_list = data.get("sec_pdg_list")
+    if sec_E_list is not None:
+        # The loader already padded sec_*_list to some k_max (see
+        # giant.data.loader.iter_file_chunks); that padded width is
+        # authoritative over whatever this call happened to pass in, so the
+        # two can never drift apart (docs/v0.3.0-design.md §8).
+        k_max = sec_E_list.shape[1]
+
+    # Clamp the classification label to k_max: the head only has k_max+1
+    # classes (0..k_max), and truncating here mirrors the k_max-slot
+    # truncation already applied to sec_cont by the loader's list padding.
+    # Without this, a rare high-multiplicity step (real data goes up to ~37)
+    # hands cross_entropy an out-of-range target and CUDA asserts.
+    n_sec = np.minimum(n_sec_raw, k_max).astype(np.int64)  # (N,)
 
     if sec_E_list is not None and sec_dir_list is not None and sec_pdg_list is not None:
-        sec_valid = np.arange(K_MAX)[None, :] < n_sec_raw[:, None]  # (N, K_MAX)
+        sec_valid = np.arange(k_max)[None, :] < n_sec_raw[:, None]  # (N, k_max)
         sec_cont = encode_secondaries(
             sec_E_list,
             sec_dir_list,
@@ -990,11 +1049,11 @@ def build_features(
             data["pre_dir"],
             sec_pdg_list=sec_pdg_list,
             phys_only=sec_phys_only,
-        )  # (N, K_MAX, 6)
+        )  # (N, k_max, 6)
         sec_type_idx = (
             encode_secondary_type_idx(sec_pdg_list, sec_valid, sec_type_class_map)
             if sec_type_class_map is not None
-            else np.zeros((len(n_sec), K_MAX), dtype=np.int64)
+            else np.zeros((len(n_sec), k_max), dtype=np.int64)
         )
     else:
         # Guard against silently training Stage 2 on zeroed targets: if any step
@@ -1017,8 +1076,8 @@ def build_features(
                 "require_secondaries=False for Stage-1-only use."
             )
         N = len(n_sec)
-        sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32)
-        sec_type_idx = np.zeros((N, K_MAX), dtype=np.int64)
+        sec_cont = np.zeros((N, k_max, 6), dtype=np.float32)
+        sec_type_idx = np.zeros((N, k_max), dtype=np.int64)
 
     if fit:
         cond_normalizer = Normalizer().fit(cond_cont)

giant/model/network.py

@@ -990,7 +990,12 @@ class Stage1Model(nn.Module):
     (docs/v0.3.0-design.md §2) moves 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`)."""
+    against (see `_migrate_legacy_model_config`).
+
+    `cond_enc`, if given, is used in place of building a fresh
+    `ConditionEncoder` — `conditioning.share_stages = true` (§3.1): `build_models`
+    constructs one shared instance and passes it to both stages, halving the
+    conditioning parameter count and forcing a common representation."""
 
     def __init__(
         self,
@@ -1008,12 +1013,17 @@ class Stage1Model(nn.Module):
         noise_dim: int = 64,
         router: Router | None = None,
         n_sec_head_k_max: int | None = None,
+        cond_enc: ConditionEncoder | None = None,
     ) -> None:
         super().__init__()
         self.generator_kind = generator
         self.noise_dim = noise_dim
-        self.cond_enc = ConditionEncoder(
-            pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
+        self.cond_enc = (
+            cond_enc
+            if cond_enc is not None
+            else ConditionEncoder(
+                pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
+            )
         )
         has_time = generator in ("flow", "ddpm")
         self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
@@ -1083,6 +1093,9 @@ class Stage2OneShot(nn.Module):
     Under `generator == "wgan"` the type slice stays folded into `sec_dim`
     (just `emb_dim` instead of `PARTICLE_PHYS_DIM` wide) and `type_head` is
     unused (`None`) — the WGAN trainer handles the ST-Gumbel relaxation.
+
+    `cond_enc`, if given, is used in place of building a fresh
+    `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`).
     """
 
     def __init__(
@@ -1105,6 +1118,7 @@ class Stage2OneShot(nn.Module):
         router: Router | None = None,
         build_n_sec_head: bool = True,
         particle_type_cfg: dict | None = None,
+        cond_enc: ConditionEncoder | None = None,
     ) -> None:
         super().__init__()
         self.generator_kind = generator
@@ -1112,8 +1126,12 @@ class Stage2OneShot(nn.Module):
         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, particle_cfg["emb_dim"])
-        self.cond_enc = ConditionEncoder(
-            pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
+        self.cond_enc = (
+            cond_enc
+            if cond_enc is not None
+            else ConditionEncoder(
+                pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
+            )
         )
         self.context_adapter = ContextAdapter(x_dim, context_dim)
         self.fuse = nn.Sequential(
@@ -1227,6 +1245,9 @@ class Stage2Autoregressive(nn.Module):
     on token position; `_token_cond` additionally fuses in the history
     encoding and two running scalars (remaining energy-budget fraction,
     normalized slot index), and feeds `forward`/`predict_type`/the trunk.
+
+    `cond_enc`, if given, is used in place of building a fresh
+    `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`).
     """
 
     def __init__(
@@ -1251,6 +1272,7 @@ class Stage2Autoregressive(nn.Module):
         history: str = "markov",
         attn_n_heads: int = 4,
         attn_n_layers: int = 2,
+        cond_enc: ConditionEncoder | None = None,
     ) -> None:
         super().__init__()
         if history not in ("markov", "attention"):
@@ -1266,8 +1288,12 @@ class Stage2Autoregressive(nn.Module):
         emb_dim = particle_cfg["emb_dim"]
         self.type_dim = stage2_type_dim(self.particle_type_cfg, emb_dim)
 
-        self.cond_enc = ConditionEncoder(
-            pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
+        self.cond_enc = (
+            cond_enc
+            if cond_enc is not None
+            else ConditionEncoder(
+                pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
+            )
         )
         self.context_adapter = ContextAdapter(x_dim, context_dim)
         self.base_fuse = nn.Sequential(
@@ -1562,6 +1588,20 @@ def _migrate_legacy_model_config(model_config: dict) -> dict:
     k_max = m.get("k_max", K_MAX)
     noise_dim = m.get("noise_dim", 64)
     router_cfg = dict(m.get("router") or {})
+    expert_hidden_dim = router_cfg.pop("expert_hidden_dim", 0)
+    expert_n_blocks = router_cfg.pop("expert_n_blocks", 0)
+    if expert_hidden_dim or expert_n_blocks:
+        raise ValueError(
+            "this checkpoint's model_config.router sets expert_hidden_dim/"
+            f"expert_n_blocks to a non-default value "
+            f"({expert_hidden_dim!r}, {expert_n_blocks!r}); v0.3.0 removed "
+            "per-expert sizing (experts always inherit the stage's "
+            "hidden_dim/n_res_blocks), so this checkpoint's routed experts "
+            "have a different width/depth than the monolith — silently "
+            "dropping these keys would resize the experts instead of "
+            "refusing (docs/v0.3.0-design.md §4.2). This checkpoint can "
+            "only be loaded by v0.2 code."
+        )
     router_cfg.setdefault("enabled", False)
 
     return {
@@ -1657,6 +1697,13 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
     instance rather than building a second, independently-parameterized one
     (v0.2's actual — probably accidental — behaviour: two routers built from
     one config with no semantic relationship between them).
+
+    `conditioning.share_stages = true` (§3.1) builds one `ConditionEncoder`
+    instance here and passes it to both stages (`Stage1Model`/`Stage2OneShot`/
+    `Stage2Autoregressive`'s `cond_enc` param), instead of each stage
+    building its own — halving the conditioning parameter count and forcing a
+    common representation. `false` (default) keeps v0.2 behaviour:
+    independent instances with identical config but independent weights.
     """
     cfg = (
         model_config
@@ -1666,17 +1713,17 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
     pdg_vocab = cfg["pdg_vocab"]
     mat_vocab = cfg["mat_vocab"]
     conditioning = cfg["conditioning"]
-    if conditioning.get("share_stages"):
-        raise NotImplementedError(
-            "conditioning.share_stages = true is not implemented yet — each "
-            "stage always builds its own ConditionEncoder for now"
-        )
     particle_cfg = conditioning["particle"]
     material_cfg = conditioning["material"]
     particle_conditioning = particle_cfg["type"]
     s1cfg = cfg["stage1_model"]
     s2cfg = cfg["stage2_model"]
     cond_out_dim = conditioning.get("out_dim", 128)
+    shared_cond_enc: ConditionEncoder | None = None
+    if conditioning.get("share_stages"):
+        shared_cond_enc = ConditionEncoder(
+            pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
+        )
 
     result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None}
 
@@ -1707,6 +1754,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
             noise_dim=(s1cfg.get("wgan") or {}).get("noise_dim", 64),
             router=stage1_router,
             n_sec_head_k_max=n_sec_head_k_max,
+            cond_enc=shared_cond_enc,
         )
 
     if s2cfg.get("active", True):
@@ -1748,6 +1796,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 history=ar_cfg.get("history", "markov"),
                 attn_n_heads=ar_cfg.get("attn_n_heads", 4),
                 attn_n_layers=ar_cfg.get("attn_n_layers", 2),
+                cond_enc=shared_cond_enc,
             )
         else:
             sec_dim = stage2_trunk_sec_dim(
@@ -1771,6 +1820,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 router=stage2_router,
                 build_n_sec_head=legacy_owner != "stage1",
                 particle_type_cfg=particle_type_cfg,
+                cond_enc=shared_cond_enc,
             )
 
     return result

giant/model/schedule.py

@@ -108,12 +108,13 @@ def flow_matching_loss_secondary(
     averaged over their own width first and then combined with equal weight
     — this stays correct if type_dim/CONT_SLOT_DIM change.
     """
-    from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM
+    from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM
 
     if type_dim is None:
         type_dim = PARTICLE_PHYS_DIM
 
     B = x1.size(0)
+    k_max = sec_mask.size(1)
     t = torch.rand(B, device=x1.device)
     x0 = torch.randn_like(x1)
     x_t = (1.0 - t.view(-1, 1)) * x0 + t.view(-1, 1) * x1
@@ -121,8 +122,8 @@ def flow_matching_loss_secondary(
     v_t = model(x_t, cond_cont, cond_cat, stage1_out, t=t)
 
     slot_dim = CONT_SLOT_DIM + type_dim
-    err = ((v_t - u_t) ** 2).view(B, K_MAX, slot_dim)
-    cont_err = err[..., :CONT_SLOT_DIM].mean(dim=-1)  # (B, K_MAX)
+    err = ((v_t - u_t) ** 2).view(B, k_max, slot_dim)
+    cont_err = err[..., :CONT_SLOT_DIM].mean(dim=-1)  # (B, k_max)
 
     mask = sec_mask.float()
     denom = mask.sum().clamp(min=1)

giant/pipeline.py

@@ -9,7 +9,6 @@ from torch.utils.data import DataLoader
 from giant import config
 from giant.constants import (
     COND_DIM,
-    K_MAX,
     PARTICLE_PHYS_DIM,
     X_DIM,
 )
@@ -106,11 +105,9 @@ def run_setup_stage(
     Stage-2 normalizers.
 
     `cfg` is the full merged v0.3 config (`conditioning`/`stage1_model`/
-    `stage2_model`), already passed through `giant.config.validate_config` —
-    in particular, `conditioning.particle.type == conditioning.material.type`
-    is assumed (the data pipeline doesn't support mixed conditioning types
-    yet, see `validate_config`), so a single shared conditioning string
-    drives `giant.data.transforms.build_features`.
+    `stage2_model`), already passed through `giant.config.validate_config`.
+    `conditioning.particle.type` and `conditioning.material.type` are
+    independent (docs/v0.3.0-design.md §3.1) and may differ.
 
     Reads from and writes to the `giant.data.setup_cache` sidecar when
     `cache_setup` is set (`rebuild_setup_cache` ignores — but still
@@ -118,7 +115,9 @@ def run_setup_stage(
     is mutated in place: an active `EnergyRouter` (`router.type == "energy"`)
     gets its `centers_init` seeded from real data quantiles here.
     """
-    conditioning = cfg["conditioning"]["particle"]["type"]
+    particle_conditioning = cfg["conditioning"]["particle"]["type"]
+    material_conditioning = cfg["conditioning"]["material"]["type"]
+    k_max = cfg["stage2_model"]["k_max"]
     stage1_router = cfg["stage1_model"].get("router") or {}
     stage2_router = cfg["stage2_model"].get("router") or {}
 
@@ -255,7 +254,9 @@ def run_setup_stage(
         for r in (stage1_router, stage2_router)
     )
     energy_idx = 3
-    norm_key = setup_cache.normalizer_key(val_fraction, seed, conditioning)
+    norm_key = setup_cache.normalizer_key(
+        val_fraction, seed, particle_conditioning, material_conditioning
+    )
     entry = cache.normalizers.get(norm_key) if cache is not None else None
 
     if entry is not None:
@@ -284,7 +285,7 @@ def run_setup_stage(
             _ReservoirSampler(capacity=100_000) if collect_energy_sample else None
         )
         for i, path in enumerate(files):
-            for chunk in iter_file_chunks(path, offset=event_id_offset(i)):
+            for chunk in iter_file_chunks(path, offset=event_id_offset(i), k_max=k_max):
                 mask = sorted_membership(chunk["event_id"], events_arr)
                 if not mask.any():
                     continue
@@ -296,15 +297,17 @@ def run_setup_stage(
                         mat_map,
                         proc_map=proc_map,
                         require_secondaries=True,
-                        conditioning=conditioning,
+                        particle_conditioning=particle_conditioning,
+                        material_conditioning=material_conditioning,
                         sec_phys_only=True,
+                        k_max=k_max,
                     )
                 )
                 cond_acc.update(cond_cont)
                 tgt_acc.update(target_s1)
                 if energy_sampler is not None:
                     energy_sampler.update(cond_cont[:, energy_idx])
-                sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None]
+                sec_valid = np.arange(sec_cont.shape[1])[None, :] < n_sec[:, None]
                 sec_phys = sec_cont[:, :, 4:6][sec_valid]
                 if len(sec_phys) > 0:
                     sec_phys_acc.update(sec_phys)
@@ -375,7 +378,9 @@ def run_train_job(
         )
 
     config.validate_config(cfg)
-    conditioning = cfg["conditioning"]["particle"]["type"]
+    particle_conditioning = cfg["conditioning"]["particle"]["type"]
+    material_conditioning = cfg["conditioning"]["material"]["type"]
+    k_max = cfg["stage2_model"]["k_max"]
     setup = run_setup_stage(
         data,
         val_fraction=t["val_fraction"],
@@ -399,18 +404,20 @@ def run_train_job(
     )
 
     # cond_cat's onehot columns (docs/v0.3.0-design.md decision 4) are
-    # present only under conditioning="onehot" — validate_config enforces
-    # conditioning.particle.type == conditioning.material.type, so both
-    # axes' topN columns are always present or absent together. run_setup_stage
-    # builds both maps whenever conditioning=="onehot" (see its own
-    # particle_cfg["type"] == "onehot" check), so they're guaranteed non-None
-    # here — asserted, not just assumed, so a future wiring bug fails loudly
-    # instead of silently dropping the onehot columns.
+    # present per-axis, independently, under that axis's own
+    # conditioning.{particle,material}.type == "onehot" (§3.1: the two axes
+    # may mix freely). run_setup_stage builds each map whenever its own axis
+    # is "onehot" (see its own particle_cfg["type"]/material_cfg["type"]
+    # checks), so they're guaranteed non-None here — asserted, not just
+    # assumed, so a future wiring bug fails loudly instead of silently
+    # dropping the onehot columns.
     cond_pdg_topn = None
     cond_mat_topn = None
-    if conditioning == "onehot":
-        assert setup.pdg_topn_map is not None and setup.mat_topn_map is not None
+    if particle_conditioning == "onehot":
+        assert setup.pdg_topn_map is not None
         cond_pdg_topn = setup.pdg_topn_map.class_map
+    if material_conditioning == "onehot":
+        assert setup.mat_topn_map is not None
         cond_mat_topn = setup.mat_topn_map.class_map
 
     # The secondary type-index map depends on stage2_model.particle_type.target,
@@ -441,11 +448,13 @@ def run_train_job(
         shuffle_buffer=shuffle_buffer,
         shuffle=True,
         proc_map=proc_map,
-        conditioning=conditioning,
+        particle_conditioning=particle_conditioning,
+        material_conditioning=material_conditioning,
         sec_phys_normalizer=sec_phys_norm,
         pdg_topn_map=cond_pdg_topn,
         mat_topn_map=cond_mat_topn,
         sec_type_class_map=sec_type_class_map,
+        k_max=k_max,
     )
     val_ds = StreamingStepsDataset(
         files=files,
@@ -457,11 +466,13 @@ def run_train_job(
         batch_size=t["batch_size"],
         shuffle=False,
         proc_map=proc_map,
-        conditioning=conditioning,
+        particle_conditioning=particle_conditioning,
+        material_conditioning=material_conditioning,
         sec_phys_normalizer=sec_phys_norm,
         pdg_topn_map=cond_pdg_topn,
         mat_topn_map=cond_mat_topn,
         sec_type_class_map=sec_type_class_map,
+        k_max=k_max,
     )
 
     pin = device.type == "cuda"

giant/rollout.py

@@ -352,7 +352,7 @@ def make_seed_frontier(
     pre_pos: np.ndarray,
     pre_E: np.ndarray,
     pre_dir: np.ndarray,
-    conditioning: str = "embedding",
+    particle_conditioning: str = "embedding",
 ) -> tuple[dict[str, np.ndarray], dict[int, int]]:
     """Build the initial frontier from primary entry states.
 
@@ -372,17 +372,17 @@ def make_seed_frontier(
     dir_ = dir_ / np.clip(np.linalg.norm(dir_, axis=1, keepdims=True), 1e-12, None)
 
     pdg_arr = np.asarray(pdg, dtype=np.int64)
-    if conditioning == "physical":
+    if particle_conditioning == "physical":
         # Real primaries always have a genuine ground-truth PDG code, looked
         # up once here and carried forward unchanged for the track's lifetime
         # (its species never changes mid-track) — same lifecycle as "pdg"
         # itself.
         mass, charge = particle_phys_array(pdg_arr).T
     else:
-        # "embedding" mode never reads mass/charge (see
+        # "embedding"/"onehot" never read mass/charge (see
         # _physical_cond_columns), so resolving them here would only risk
-        # crashing an embedding-mode rollout on a PDG code giant.particles
-        # can't resolve, for a value that's never used.
+        # crashing a rollout on a PDG code giant.particles can't resolve, for
+        # a value that's never used.
         mass = np.zeros(n, dtype=np.float64)
         charge = np.zeros(n, dtype=np.float64)
 
@@ -457,8 +457,10 @@ def rollout(
     max_tracks_per_event: int | None = None,
     escape_threshold: float | None = None,
     on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None,
-    conditioning: str = "embedding",
+    particle_conditioning: str = "embedding",
+    material_conditioning: str = "embedding",
     pdg_topn_map: "TopNMap | None" = None,
+    mat_topn_map: "TopNMap | None" = None,
     other_policy: str = "sample",
     seed: int | None = None,
     stage1_ddpm_steps: int = 1000,
@@ -485,8 +487,13 @@ def rollout(
     autoregressive) is inferred from `sec_decoder`'s own class — see
     `sample_stage1`/`sample_stage2` (giant.sample).
 
-    `pdg_topn_map`/`other_policy` are only read under
-    `stage2_model.particle_type.target = "onehot"` (§3.3); `seed` seeds the
+    `pdg_topn_map`/`mat_topn_map` serve two independent purposes that happen
+    to share `pdg_topn_map` (docs/v0.3.0-design.md §8 — one PDG map, not
+    two): they're required whenever `particle_conditioning`/
+    `material_conditioning` is `"onehot"` (feeds `build_cond_features`'s
+    extra `cond_cat` top-N columns — §3.1), and `pdg_topn_map`/`other_policy`
+    are additionally read under `stage2_model.particle_type.target =
+    "onehot"` (§3.3, secondary-species decode). `seed` seeds the
     `other_policy = "sample"` draw only (torch/numpy sampling itself is
     seeded by the caller, same as today).
 
@@ -494,6 +501,16 @@ def rollout(
     diagnostic across the whole run — see `L1DistCollector`. Only populated
     under `particle_type.target = "embedding"`; a no-op otherwise.
     """
+    if particle_conditioning == "onehot" and pdg_topn_map is None:
+        raise RuntimeError(
+            "conditioning.particle.type='onehot' rollout needs pdg_topn_map "
+            "(the checkpoint's saved top-N map) — see ckpt['pdg_topn_map']"
+        )
+    if material_conditioning == "onehot" and mat_topn_map is None:
+        raise RuntimeError(
+            "conditioning.material.type='onehot' rollout needs mat_topn_map "
+            "(the checkpoint's saved top-N map) — see ckpt['mat_topn_map']"
+        )
     device = device or torch.device("cpu")
     stage1_model.eval()
     sec_decoder.eval()
@@ -507,7 +524,7 @@ def rollout(
         seeds["pre_pos"],
         seeds["pre_E"],
         seeds["pre_dir"],
-        conditioning=conditioning,
+        particle_conditioning=particle_conditioning,
     )
     rec = _Recorder(sink=on_chunk)
 
@@ -534,8 +551,10 @@ def rollout(
                     steps,
                     device,
                     max_tracks_per_event,
-                    conditioning,
+                    particle_conditioning,
+                    material_conditioning,
                     pdg_topn_map,
+                    mat_topn_map,
                     other_policy,
                     rng,
                     stage1_ddpm_steps,
@@ -570,8 +589,10 @@ def _step_chunk(
     steps,
     device,
     max_tracks_per_event,
-    conditioning,
+    particle_conditioning,
+    material_conditioning,
     pdg_topn_map,
+    mat_topn_map,
     other_policy,
     rng,
     stage1_ddpm_steps,
@@ -587,7 +608,7 @@ def _step_chunk(
     tr["_material"] = material
     tr["_layer_id"] = layer_id
 
-    if conditioning == "physical":
+    if particle_conditioning == "physical":
         # Under physical-property conditioning, mass/charge (already resolved
         # on every track — see the cond_dict comment below) drive the model,
         # not a training-vocab PDG embedding — build_cond_features passes
@@ -655,7 +676,18 @@ def _step_chunk(
         "charge": tr["charge"],
     }
     cond_cont, cond_cat = build_cond_features(
-        cond_dict, pdg_map, mat_map, cond_norm, conditioning=conditioning
+        cond_dict,
+        pdg_map,
+        mat_map,
+        cond_norm,
+        particle_conditioning=particle_conditioning,
+        material_conditioning=material_conditioning,
+        pdg_topn_map=pdg_topn_map.class_map
+        if particle_conditioning == "onehot"
+        else None,
+        mat_topn_map=mat_topn_map.class_map
+        if material_conditioning == "onehot"
+        else None,
     )
     cc = torch.from_numpy(cond_cont).float().to(device)
     ck = torch.from_numpy(cond_cat).long().to(device)

giant/train.py

@@ -4,7 +4,6 @@ import math
 import os
 import signal
 import time
-import warnings
 from pathlib import Path
 from types import FrameType
 from typing import Callable
@@ -16,7 +15,7 @@ import torch.optim as optim
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
-from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM
+from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM
 from giant.data.loader import TopNMap
 from giant.data.setup_cache import topnmap_to_json
 from giant.model.network import Router, stage2_type_dim
@@ -483,9 +482,10 @@ class FlowDDPMStageTrainer(StageTrainer):
     ) -> None:
         if is_stage2 and generator not in ("flow",):
             raise NotImplementedError(
-                f"stage2_model.generator={generator!r} is not implemented for "
-                "stage 2 (only 'flow' and 'wgan' — 'ddpm' has no stage-2 "
-                "secondary-decoder loss yet, see docs/v0.3.0-design.md)"
+                f"stage2_model.generator={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 — "
+                "see docs/v0.3.0-design.md §11.2)"
             )
         self.name = name
         self.is_stage2 = is_stage2
@@ -653,7 +653,9 @@ class FlowDDPMStageTrainer(StageTrainer):
             proc_idx,
             sec_type_idx,
         ) = _batch_to_device(batch, device)
-        sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1)
+        sec_mask = torch.arange(sec_cont.size(1), device=device).unsqueeze(
+            0
+        ) < n_sec.unsqueeze(1)
         stage1_ctx = x1_s1.detach()
 
         x1_s2 = None
@@ -957,8 +959,9 @@ class WGANStageTrainer(StageTrainer):
                 self.particle_type_cfg, self.particle_type_emb_dim
             )
             slot_width = CONT_SLOT_DIM + type_dim
+            k_max = sec_cont.size(1)
 
-            sec_mask = torch.arange(K_MAX, device=device).unsqueeze(
+            sec_mask = torch.arange(k_max, device=device).unsqueeze(
                 0
             ) < n_sec.unsqueeze(1)
             mask = (
@@ -1002,7 +1005,7 @@ class WGANStageTrainer(StageTrainer):
                     ).reshape(B, -1)
                     * mask
                 )
-                z = torch.randn(B, K_MAX, self.model.noise_dim, device=device)
+                z = torch.randn(B, k_max, self.model.noise_dim, device=device)
                 fake_raw = self.model(
                     z,
                     cond_cont,
@@ -1042,7 +1045,7 @@ class WGANStageTrainer(StageTrainer):
                     self.type_gumbel_tau_end,
                 )
                 fake_raw = _relax_onehot_type_slice(
-                    fake_raw, K_MAX, CONT_SLOT_DIM, type_dim, tau, grad_probe=grad_probe
+                    fake_raw, k_max, CONT_SLOT_DIM, type_dim, tau, grad_probe=grad_probe
                 )
             fake = fake_raw * mask
 
@@ -1161,39 +1164,14 @@ class WGANStageTrainer(StageTrainer):
             group["lr"] = resumed_critic_lr
 
 
-_WARNED_MARGINAL_VALIDATION_BROKEN = False
-
-
 def _try_validate_marginals(trainer: StageTrainer, val_loader, device, **kwargs):
-    """`giant.sample` itself has a per-stage generator/n_sec-ownership
-    update (v0.3.0 step 6 — `sample_stage1`/`sample_stage2`/`resolve_n_sec`,
-    docs/v0.3.0-design.md §10); `giant/validate.py` hasn't been updated to
-    call it yet, and also still unpacks the val batch as a 6-tuple, which no
-    longer matches `StreamingStepsDataset`'s 7-tuple shape (the `sec_type_idx`
-    column step 4 added). Both are `giant/validate.py`-side gaps, not
-    `giant/sample.py`'s — tracked, not yet scheduled. Until fixed, this
-    degrades gracefully with a one-time warning instead of crashing the
-    whole training run.
-    """
-    global _WARNED_MARGINAL_VALIDATION_BROKEN
+    """Runs `validate_marginals` on `trainer`'s sampling model (EMA model if
+    present, else the raw model). `validate_marginals` itself dispatches
+    through `giant.sample.sample_stage1`/`sample_stage2`/`resolve_n_sec`
+    (docs/v0.3.0-design.md §10), so this is generator- and
+    one-shot-vs-autoregressive-agnostic."""
     model = trainer.sampling_model()
-    try:
-        return validate_marginals(model, val_loader, device=device, **kwargs)
-    except Exception as exc:  # noqa: BLE001 — see docstring: any failure here
-        # is expected until giant/validate.py is updated for the v0.3.0
-        # per-stage dispatch + 7-tuple batch shape.
-        if not _WARNED_MARGINAL_VALIDATION_BROKEN:
-            warnings.warn(
-                "marginal validation unavailable this run "
-                f"({type(exc).__name__}: {exc}) — giant/validate.py hasn't "
-                "been updated for v0.3.0's per-stage sample dispatch / "
-                "7-tuple val batch shape yet; val_marginal_kl stays NaN, and "
-                "wgan best-checkpoint selection falls back to the "
-                "Wasserstein-distance magnitude.",
-                stacklevel=2,
-            )
-            _WARNED_MARGINAL_VALIDATION_BROKEN = True
-        return None
+    return validate_marginals(model, val_loader, device=device, **kwargs)
 
 
 def _build_stage_trainers(
@@ -1688,7 +1666,6 @@ def train(
                         stage1_tr,
                         val_loader,
                         device,
-                        mode=cfg["stage1_model"]["generator"],
                         sec_decoder=trainers["stage2"].sampling_model()
                         if "stage2" in trainers
                         else None,
@@ -1716,11 +1693,11 @@ def train(
                         stage1_tr,
                         val_loader,
                         device,
-                        mode=cfg["stage1_model"]["generator"],
-                        schedule=stage1_tr.ddpm_schedule
-                        if isinstance(stage1_tr, FlowDDPMStageTrainer)
-                        else None,
                         steps=validate_steps,
+                        ddpm_steps=stage1_tr.ddpm_schedule.T
+                        if isinstance(stage1_tr, FlowDDPMStageTrainer)
+                        and stage1_tr.ddpm_schedule is not None
+                        else 1000,
                         sec_decoder=trainers["stage2"].sampling_model()
                         if "stage2" in trainers
                         else None,

giant/validate.py

@@ -2,23 +2,12 @@ import numpy as np
 import torch
 from torch.utils.data import DataLoader
 
-from giant.constants import K_MAX, LOCAL_TARGET_NAMES
-from giant.sample import (
-    sample_flow,
-    sample_ddpm,
-    sample_ddim,
-    sample_secondaries,
-    sample_wgan,
-    sample_secondaries_wgan,
-)
+from giant.constants import LOCAL_TARGET_NAMES
+from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
 
 _SEC_PHYS_NAMES = ["log_mass", "charge"]
 
 
-def _kw(steps: int | None) -> dict[str, int]:
-    return {} if steps is None else {"steps": steps}
-
-
 def _histogram_kl(
     p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8
 ) -> float:
@@ -44,15 +33,44 @@ def _bincount_frac(x: np.ndarray, minlength: int) -> np.ndarray:
     return counts / total if total > 0 else counts
 
 
+def _categorical_kl(
+    real_idx: np.ndarray, gen_idx: np.ndarray, n_classes: int, eps: float = 1e-8
+) -> float:
+    """KL(P_real || Q_gen) between two class-index samples over `n_classes`
+    categories, estimated from bincount fractions. NaN if either side has no
+    valid samples (mirrors `_histogram_kl`'s empty-input handling)."""
+    if len(real_idx) == 0 or len(gen_idx) == 0:
+        return float("nan")
+    p = _bincount_frac(real_idx, n_classes) + eps
+    q = _bincount_frac(gen_idx, n_classes) + eps
+    p /= p.sum()
+    q /= q.sum()
+    return float(np.sum(p * np.log(p / q)))
+
+
+def _embedding_nearest_class(
+    vectors: torch.Tensor, emb_weight: torch.Tensor
+) -> np.ndarray:
+    """Nearest row index (L1) of `vectors` (..., emb_dim) against `emb_weight`
+    (vocab, emb_dim) — same computation as
+    `giant.particles.decode_embedding_nearest`, but returning the raw class
+    index instead of a decoded PDG code: validate.py only needs a
+    real-vs-generated class-distribution comparison, not a rollout-usable
+    identity, so there's no need for the pdg_map inversion here."""
+    flat = vectors.reshape(-1, vectors.size(-1))
+    dist = (flat.unsqueeze(1) - emb_weight.detach().unsqueeze(0)).abs().sum(-1)
+    nearest = dist.argmin(dim=1)
+    return nearest.reshape(vectors.shape[:-1]).cpu().numpy()
+
+
 def validate_marginals(
-    model: torch.nn.Module,
+    stage1_model: torch.nn.Module,
     val_loader: DataLoader,
-    mode: str = "flow",
-    schedule=None,
     device: torch.device | None = None,
     n_batches: int | None = None,
     kl_bins: int = 50,
-    steps: int | None = None,
+    steps: int = 10,
+    ddpm_steps: int = 1000,
     sec_decoder: torch.nn.Module | None = None,
 ) -> dict[str, np.ndarray | float]:
     """Compare per-dimension marginals of generated vs. real steps.
@@ -61,51 +79,61 @@ def validate_marginals(
     normalised space. `kl_divergence[j]` is KL(real || generated) for
     dimension j, estimated from a shared histogram over both samples.
 
-    `steps` overrides the number of sampler steps (flow ODE steps or DDIM
-    substeps); `None` keeps each sampler's own default. Unused in "ddpm"
-    mode, which always runs the full schedule.
+    Stage 1 is sampled via `giant.sample.sample_stage1`, which dispatches on
+    `stage1_model.generator_kind` — `steps`/`ddpm_steps` are forwarded but
+    only one of them is actually read, depending on that dispatch.
 
-    When `sec_decoder` is given, also validates Stage 2: n_sec distribution
-    (+ classification accuracy), predicted secondary physical-identity
-    (log_mass, charge) marginals, and per-slot energy-fraction marginals —
-    restricted to each side's own valid slots (real: `n_sec`; generated: the
-    Stage-1 head's argmax), since the two need not agree on how many slots
-    are valid. Compared directly in normalised space (no denormalising —
-    KL estimated from a shared per-sample histogram is invariant to a shared
-    affine rescaling of both sides). Adds {"n_sec_real", "n_sec_pred",
-    "n_sec_accuracy", "phys_real", "phys_generated", "phys_kl",
-    "energy_fraction_kl"} to the returned dict.
+    When `sec_decoder` is given, also validates Stage 2 via
+    `giant.sample.sample_stage2`/`resolve_n_sec` (generator- and
+    one-shot-vs-autoregressive-agnostic, docs/v0.3.0-design.md §10): n_sec
+    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
+    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
+    `target = "physical"`, {"phys_real", "phys_generated", "phys_kl"}
+    (continuous log_mass/charge marginals — v0.2 behaviour), or under
+    `target` in `("onehot", "embedding")`, {"type_class_real",
+    "type_class_gen", "type_class_kl"} (categorical class-index marginal:
+    argmax for "onehot", L1-nearest conditioning-embedding row for
+    "embedding" — see `_embedding_nearest_class`). Compared directly in
+    normalised space (no denormalising — KL estimated from a shared
+    per-sample histogram/bincount is invariant to a shared affine rescaling
+    of both sides).
     """
     if device is None:
-        device = next(model.parameters()).device
-    model.eval()
+        device = next(stage1_model.parameters()).device
+    stage1_model.eval()
     if sec_decoder is not None:
         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"
+    )
+
     all_real, all_gen = [], []
     all_n_sec_real, all_n_sec_pred = [], []
     all_phys_real, all_phys_gen = [], []
-    all_frac_real: list[list[np.ndarray]] = [[] for _ in range(K_MAX)]
-    all_frac_gen: list[list[np.ndarray]] = [[] for _ in range(K_MAX)]
+    all_type_class_real, all_type_class_gen = [], []
+    all_frac_real: list[list[np.ndarray]] = [[] for _ in range(k_max)]
+    all_frac_gen: list[list[np.ndarray]] = [[] for _ in range(k_max)]
 
     for i, batch in enumerate(val_loader):
         if n_batches is not None and i >= n_batches:
             break
-        # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx).
-        cond_cont, cond_cat, x1, n_sec, sec_cont, _proc_idx = batch
+        # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx,
+        # sec_type_idx).
+        cond_cont, cond_cat, x1, n_sec, sec_cont, _proc_idx, sec_type_idx = batch
         cond_cont = cond_cont.to(device)
         cond_cat = cond_cat.to(device)
 
-        if mode == "flow":
-            gen, n_sec_pred = sample_flow(model, cond_cont, cond_cat, **_kw(steps))
-        elif mode == "ddpm":
-            gen, n_sec_pred = sample_ddpm(model, cond_cont, cond_cat, schedule)
-        elif mode == "wgan":
-            gen, n_sec_pred = sample_wgan(model, cond_cont, cond_cat)
-        else:
-            gen, n_sec_pred = sample_ddim(
-                model, cond_cont, cond_cat, schedule, **_kw(steps)
-            )
+        gen, n_sec_pred = sample_stage1(
+            stage1_model, cond_cont, cond_cat, steps=steps, ddpm_steps=ddpm_steps
+        )
 
         all_real.append(x1.numpy())
         all_gen.append(gen.cpu().numpy())
@@ -113,37 +141,41 @@ def validate_marginals(
         if sec_decoder is None:
             continue
 
+        n_sec_pred = resolve_n_sec(
+            stage1_model, sec_decoder, cond_cont, cond_cat, gen, n_sec_pred
+        )
         n_sec_pred_np = n_sec_pred.cpu().numpy()
         n_sec_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_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)))
-        real_phys = sec_cont[:, :, 4:6].numpy()  # (B, K_MAX, 2) [log_mass, charge]
 
-        if mode == "wgan":
-            sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries_wgan(
-                sec_decoder, cond_cont, cond_cat, gen, n_sec_pred
-            )
-        else:
-            sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries(
-                sec_decoder,
-                cond_cont,
-                cond_cat,
-                gen,
-                n_sec_pred,
-                steps=steps if steps is not None else 10,
-            )
+        sec_cont_pred, sec_type_pred, sec_valid_pred = sample_stage2(
+            sec_decoder, cond_cont, cond_cat, gen, n_sec_pred, steps=steps
+        )
         gen_frac = 1.0 / (
             1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64))
         )
-        gen_phys = sec_phys_pred.cpu().numpy()
         gen_valid = sec_valid_pred.cpu().numpy()
 
-        all_phys_real.append(real_phys[real_valid])
-        all_phys_gen.append(gen_phys[gen_valid])
-        for j in range(K_MAX):
+        if target == "physical":
+            real_phys = sec_cont[:, :, 4:6].numpy()  # (B, k_max, 2) [log_mass, charge]
+            gen_phys = sec_type_pred.cpu().numpy()
+            all_phys_real.append(real_phys[real_valid])
+            all_phys_gen.append(gen_phys[gen_valid])
+        else:
+            sec_type_idx_np = sec_type_idx.numpy()
+            all_type_class_real.append(sec_type_idx_np[real_valid])
+            if target == "onehot":
+                gen_class = sec_type_pred.argmax(dim=-1).cpu().numpy()
+            else:  # "embedding"
+                emb_weight = sec_decoder.cond_enc.pdg_emb.weight
+                gen_class = _embedding_nearest_class(sec_type_pred, emb_weight)
+            all_type_class_gen.append(gen_class[gen_valid])
+
+        for j in range(k_max):
             all_frac_real[j].append(real_frac[real_valid[:, j], j])
             all_frac_gen[j].append(gen_frac[gen_valid[:, j], j])
 
@@ -181,20 +213,8 @@ def validate_marginals(
     n_sec_real = np.concatenate(all_n_sec_real, axis=0)
     n_sec_pred_all = np.concatenate(all_n_sec_pred, axis=0)
     n_sec_accuracy = float((n_sec_real == n_sec_pred_all).mean())
-    phys_real = np.concatenate(all_phys_real, axis=0)  # (M, 2)
-    phys_gen = np.concatenate(all_phys_gen, axis=0)  # (M, 2)
 
-    if len(phys_real) > 0 and len(phys_gen) > 0:
-        phys_kl = np.array(
-            [
-                _histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins)
-                for j in range(2)
-            ]
-        )
-    else:
-        phys_kl = np.full(2, np.nan)
-
-    energy_fraction_kl = np.full(K_MAX, np.nan)
+    energy_fraction_kl = np.full(k_max, np.nan)
     print(
         f"\n{'n_sec':<20} accuracy={n_sec_accuracy:.4f}  "
         f"mean|Δ|={np.abs(n_sec_real - n_sec_pred_all).mean():.4f}"
@@ -208,26 +228,12 @@ def validate_marginals(
     for v in range(max_n_sec):
         print(f"{v:<20} {real_n_sec_dist[v]:>10.4f} {gen_n_sec_dist[v]:>10.4f}")
 
-    print(
-        f"\n{'sec phys (normalised)':<20} {'real_mean':>10} {'gen_mean':>10} "
-        f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
-    )
-    print("-" * 68)
-    for j, name in enumerate(_SEC_PHYS_NAMES):
-        r, g = phys_real[:, j], phys_gen[:, j]
-        if len(r) == 0 or len(g) == 0:
-            continue
-        print(
-            f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} "
-            f"{r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"
-        )
-
     print(
         f"\n{'sec slot (energy frac.)':<24} {'real_mean':>10} {'gen_mean':>10} "
         f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
     )
     print("-" * 90)
-    for j in range(K_MAX):
+    for j in range(k_max):
         r = np.concatenate(all_frac_real[j]) if all_frac_real[j] else np.array([])
         g = np.concatenate(all_frac_gen[j]) if all_frac_gen[j] else np.array([])
         if len(r) == 0 or len(g) == 0:
@@ -244,10 +250,63 @@ def validate_marginals(
             "n_sec_real": n_sec_real,
             "n_sec_pred": n_sec_pred_all,
             "n_sec_accuracy": n_sec_accuracy,
-            "phys_real": phys_real,
-            "phys_generated": phys_gen,
-            "phys_kl": phys_kl,
             "energy_fraction_kl": energy_fraction_kl,
         }
     )
+
+    if target == "physical":
+        phys_real = np.concatenate(all_phys_real, axis=0)  # (M, 2)
+        phys_gen = np.concatenate(all_phys_gen, axis=0)  # (M, 2)
+
+        if len(phys_real) > 0 and len(phys_gen) > 0:
+            phys_kl = np.array(
+                [
+                    _histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins)
+                    for j in range(2)
+                ]
+            )
+        else:
+            phys_kl = np.full(2, np.nan)
+
+        print(
+            f"\n{'sec phys (normalised)':<20} {'real_mean':>10} {'gen_mean':>10} "
+            f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
+        )
+        print("-" * 68)
+        for j, name in enumerate(_SEC_PHYS_NAMES):
+            r, g = phys_real[:, j], phys_gen[:, j]
+            if len(r) == 0 or len(g) == 0:
+                continue
+            print(
+                f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} "
+                f"{r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"
+            )
+
+        result.update(
+            {"phys_real": phys_real, "phys_generated": phys_gen, "phys_kl": phys_kl}
+        )
+    else:
+        type_class_real = np.concatenate(all_type_class_real, axis=0)
+        type_class_gen = np.concatenate(all_type_class_gen, axis=0)
+        n_classes = (
+            sec_decoder.type_dim
+            if target == "onehot"
+            else sec_decoder.cond_enc.pdg_emb.weight.size(0)
+        )
+        type_class_kl = _categorical_kl(type_class_real, type_class_gen, n_classes)
+
+        print(
+            f"\n{'sec type class (' + target + ')':<24} "
+            f"n={len(type_class_real)}/{len(type_class_gen)} "
+            f"KL(real||gen)={type_class_kl:.4f}"
+        )
+
+        result.update(
+            {
+                "type_class_real": type_class_real,
+                "type_class_gen": type_class_gen,
+                "type_class_kl": type_class_kl,
+            }
+        )
+
     return result

scripts/dwarf.py

@@ -522,10 +522,21 @@ def warm_cache(
             "--seed", "-s", help="Must match the `giant train` run(s) to warm for"
         ),
     ] = 0,
-    conditioning: Annotated[
+    particle_conditioning: Annotated[
         Conditioning,
         typer.Option(
-            "--conditioning", help="Must match the `giant train` run(s) to warm for"
+            "--particle-conditioning",
+            help="Must match the `giant train` run(s)' conditioning.particle.type "
+            "to warm for",
+        ),
+    ] = Conditioning.physical,
+    material_conditioning: Annotated[
+        Conditioning,
+        typer.Option(
+            "--material-conditioning",
+            help="Must match the `giant train` run(s)' conditioning.material.type "
+            "to warm for — independent of --particle-conditioning "
+            "(docs/v0.3.0-design.md §3.1: the two axes may differ)",
         ),
     ] = Conditioning.physical,
     router: Annotated[
@@ -552,15 +563,17 @@ 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/--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.
+    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.
     """
     run_warm_setup_cache(
         data=str(data),
         val_fraction=val_fraction,
         seed=seed,
-        conditioning=conditioning.value,
+        particle_conditioning=particle_conditioning.value,
+        material_conditioning=material_conditioning.value,
         router_enabled=router,
         router_type=router_type,
         n_experts=n_experts,

scripts/warm_setup_cache.py

@@ -9,6 +9,7 @@ for the sidecar itself.
 
 from pathlib import Path
 
+from giant.constants import K_MAX
 from giant.pipeline import run_setup_stage
 
 
@@ -16,7 +17,8 @@ def run_warm_setup_cache(
     data: str,
     val_fraction: float = 0.1,
     seed: int = 0,
-    conditioning: str = "physical",
+    particle_conditioning: str = "physical",
+    material_conditioning: str = "physical",
     router_enabled: bool = False,
     router_type: str = "energy",
     n_experts: int = 4,
@@ -25,10 +27,12 @@ def run_warm_setup_cache(
 ) -> None:
     """Populate (or refresh) the setup cache sidecar for `data`.
 
-    `val_fraction`/`seed`/`conditioning` select the normalizer cache entry
+    `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.
-    `router_enabled`/`router_type`/`n_experts` only matter for
+    `giant train` invocation will use so it hits this warmed entry. The two
+    conditioning axes are independent (docs/v0.3.0-design.md §3.1) and may
+    differ. `router_enabled`/`router_type`/`n_experts` only matter for
     `router_type == "process"` (warms that `n_experts`'s process map); the
     energy-router quantile summary is always collected regardless, so a
     later `--router-type energy` run never needs to rescan just to seed
@@ -40,16 +44,17 @@ def run_warm_setup_cache(
         "n_experts": n_experts,
     }
     # A minimal v0.3 cfg — only the keys run_setup_stage actually reads
-    # (conditioning.particle.type, stage{1,2}_model.router). This CLI only
-    # ever configures one router (matching today's single --router-type
-    # flag), so it's placed on stage1_model; stage2_model's stays disabled.
+    # (conditioning.{particle,material}.type, stage{1,2}_model.router). This
+    # CLI only ever configures one router (matching today's single
+    # --router-type flag), so it's placed on stage1_model; stage2_model's
+    # stays disabled.
     cfg = {
         "conditioning": {
-            "particle": {"type": conditioning},
-            "material": {"type": conditioning},
+            "particle": {"type": particle_conditioning},
+            "material": {"type": material_conditioning},
         },
         "stage1_model": {"router": router_cfg},
-        "stage2_model": {"router": {"enabled": False}},
+        "stage2_model": {"router": {"enabled": False}, "k_max": K_MAX},
     }
     run_setup_stage(
         Path(data),

tests/test_cli_predict.py

@@ -1,73 +1,14 @@
 import uuid
 
-import pytest
-import typer
 import yaml
 
 from giant.cli import (
     _CEPH_PREDICTIONS,
-    _check_conditioning_onehot_support,
     _resolve_prediction_output,
     _write_prediction_ref,
 )
 
 
-# ---------------------------------------------------------------------------
-# _check_conditioning_onehot_support
-# ---------------------------------------------------------------------------
-
-
-def _nested_model_cfg(
-    particle_type="physical", material_type="physical", target="physical"
-):
-    return {
-        "conditioning": {
-            "particle": {"type": particle_type, "emb_dim": 8},
-            "material": {"type": material_type, "emb_dim": 8},
-        },
-        "stage2_model": {"particle_type": {"target": target}},
-    }
-
-
-def test_check_conditioning_onehot_support_allows_physical():
-    _check_conditioning_onehot_support(_nested_model_cfg(), "predict")  # no raise
-
-
-def test_check_conditioning_onehot_support_rejects_onehot_particle_conditioning():
-    cfg = _nested_model_cfg(particle_type="onehot")
-    with pytest.raises(typer.Exit):
-        _check_conditioning_onehot_support(cfg, "predict")
-
-
-def test_check_conditioning_onehot_support_rejects_onehot_material_conditioning():
-    cfg = _nested_model_cfg(material_type="onehot")
-    with pytest.raises(typer.Exit):
-        _check_conditioning_onehot_support(cfg, "rollout")
-
-
-def test_check_conditioning_onehot_support_allows_onehot_particle_type_target():
-    """stage2_model.particle_type.target="onehot" is implemented (v0.3.0
-    step 6, giant.rollout.decode_secondary_identity) — it's a separate axis
-    from conditioning.particle.type, which this guard doesn't gate at all."""
-    cfg = _nested_model_cfg(target="onehot")
-    _check_conditioning_onehot_support(cfg, "predict")  # no raise
-
-
-def test_check_conditioning_onehot_support_allows_embedding_particle_type_target():
-    cfg = _nested_model_cfg(
-        particle_type="embedding", material_type="embedding", target="embedding"
-    )
-    _check_conditioning_onehot_support(cfg, "predict")  # no raise
-
-
-def test_check_conditioning_onehot_support_is_noop_for_v02_flat_model_config():
-    """A v0.2 checkpoint's flat model_config has conditioning as a plain
-    string, not a dict — never onehot, so this must be a silent no-op rather
-    than crash on `.get("particle")` against a string."""
-    cfg = {"conditioning": "embedding", "mode": "flow"}
-    _check_conditioning_onehot_support(cfg, "predict")  # no raise
-
-
 # ---------------------------------------------------------------------------
 # _resolve_prediction_output
 # ---------------------------------------------------------------------------

tests/test_cli_train_overrides.py

@@ -0,0 +1,97 @@
+"""Tests for `giant train`'s stage-prefixed CLI flags (docs/v0.3.0-design.md
+decision 7 / docs/v0.3.0-followups.md item 2): --stage1-*/--stage2-* must
+independently override each stage's config block, and must take precedence
+over the older shared flags (--mode/--hidden-dim/--n-critic/... ) that still
+apply the same value to both stages for backward compatibility."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from typer.testing import CliRunner
+
+import giant.cli as cli
+
+runner = CliRunner()
+
+
+def _invoke_and_capture_cfg(monkeypatch, tmp_path: Path, args: list[str]) -> dict:
+    captured: dict = {}
+
+    def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
+        captured["cfg"] = cfg
+
+    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+
+    result = runner.invoke(
+        cli.app,
+        ["train", "dummy.parquet", "--out", str(tmp_path / "run")] + args,
+    )
+    assert result.exit_code == 0, result.output
+    return captured["cfg"]
+
+
+def test_stage_prefixed_generator_overrides_shared_mode(monkeypatch, tmp_path):
+    cfg = _invoke_and_capture_cfg(
+        monkeypatch,
+        tmp_path,
+        ["--mode", "wgan", "--stage1-generator", "flow"],
+    )
+    assert cfg["stage1_model"]["generator"] == "flow"
+    assert cfg["stage2_model"]["generator"] == "wgan"
+
+
+def test_stage2_only_knobs(monkeypatch, tmp_path):
+    cfg = _invoke_and_capture_cfg(
+        monkeypatch,
+        tmp_path,
+        [
+            "--stage2-decoder",
+            "one_shot",
+            "--stage2-k-max",
+            "8",
+            "--stage2-hidden-dim",
+            "32",
+            "--stage2-context-dim",
+            "16",
+            "--stage2-stage1-context",
+            "sampled",
+        ],
+    )
+    assert cfg["stage2_model"]["decoder"] == "one_shot"
+    assert cfg["stage2_model"]["k_max"] == 8
+    assert cfg["stage2_model"]["hidden_dim"] == 32
+    assert cfg["stage2_model"]["context_dim"] == 16
+    assert cfg["stage2_model"]["stage1_context"] == "sampled"
+    # untouched stage1 defaults
+    assert cfg["stage1_model"]["hidden_dim"] == 256
+
+
+def test_stage1_hidden_dim_flag_overrides_legacy_hidden_dim_flag(monkeypatch, tmp_path):
+    cfg = _invoke_and_capture_cfg(
+        monkeypatch,
+        tmp_path,
+        ["--hidden-dim", "64", "--stage1-hidden-dim", "128"],
+    )
+    assert cfg["stage1_model"]["hidden_dim"] == 128
+
+
+def test_wgan_knobs_split_per_stage(monkeypatch, tmp_path):
+    cfg = _invoke_and_capture_cfg(
+        monkeypatch,
+        tmp_path,
+        [
+            "--mode",
+            "wgan",
+            "--n-critic",
+            "5",
+            "--stage1-n-critic",
+            "3",
+            "--stage2-gp-weight",
+            "2.5",
+        ],
+    )
+    assert cfg["stage1_model"]["wgan"]["n_critic"] == 3
+    assert cfg["stage1_model"]["wgan"]["gp_weight"] == 10.0
+    assert cfg["stage2_model"]["wgan"]["n_critic"] == 5
+    assert cfg["stage2_model"]["wgan"]["gp_weight"] == 2.5

tests/test_config.py

@@ -584,21 +584,19 @@ def test_validate_config_embedding_target_passes_with_embedding_conditioning():
     gconfig.validate_config(cfg)  # must not raise
 
 
-def test_validate_config_mixed_particle_material_conditioning_not_supported():
-    """The data pipeline doesn't support mixed conditioning types yet, even
-    though ConditionEncoder itself already can (docs/v0.3.0-design.md §3.1
-    vs. giant/data/transforms.py's still-single conditioning param)."""
+def test_validate_config_mixed_particle_material_conditioning_is_valid():
+    """docs/v0.3.0-design.md §3.1: the particle and material conditioning
+    axes are configured independently and may mix freely — e.g. material
+    "physical" with particle "embedding" — and the data pipeline
+    (giant/data/transforms.py) now implements that end-to-end, so
+    validate_config must not reject it."""
     cfg = _cfg_with(
         **{
             "conditioning.particle.type": "physical",
             "conditioning.material.type": "embedding",
         }
     )
-    try:
-        gconfig.validate_config(cfg)
-        assert False, "expected ValueError"
-    except ValueError as e:
-        assert "mixed" in str(e) or "conditioning.material.type" in str(e)
+    gconfig.validate_config(cfg)  # must not raise
 
 
 def test_validate_config_pdg_router_incompatible_with_physical_conditioning():
@@ -639,6 +637,49 @@ def test_validate_config_stop_token_not_implemented():
         assert "stop_token" in str(e)
 
 
+def test_validate_config_n_sec_truth_rejected_for_rollout_capable_checkpoint():
+    """docs/v0.3.0-design.md §9: 'n_sec.mode = "truth" is invalid for a
+    rollout-capable checkpoint' — both stages active means giant rollout
+    could load this checkpoint, but 'truth' has no ground truth to draw
+    n_sec from at rollout time."""
+    cfg = _cfg_with(
+        **{
+            "stage2_model.n_sec.mode": "truth",
+            "stage1_model.active": True,
+            "stage2_model.active": True,
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "n_sec.mode" in str(e) and "truth" in str(e)
+
+
+def test_validate_config_n_sec_truth_allowed_for_stage2_only_checkpoint():
+    """'truth' is exactly the standalone stage-2 evaluation mode the design
+    doc carves out — stage1_model.active = false must still pass."""
+    cfg = _cfg_with(
+        **{
+            "stage2_model.n_sec.mode": "truth",
+            "stage1_model.active": False,
+            "stage2_model.active": True,
+        }
+    )
+    gconfig.validate_config(cfg)  # must not raise
+
+
+def test_validate_config_n_sec_truth_allowed_when_stage2_inactive():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.n_sec.mode": "truth",
+            "stage1_model.active": True,
+            "stage2_model.active": False,
+        }
+    )
+    gconfig.validate_config(cfg)  # must not raise
+
+
 def test_validate_config_ar_default_markov_always_passes():
     """DEFAULT_CONFIG already has decoder='autoregressive',
     history='markov', teacher_forcing='always' — must not raise (v0.3.0

tests/test_dataset.py

@@ -126,7 +126,8 @@ def test_streaming_dataset_offsets_colliding_event_ids_across_files(tmp_path):
             target_normalizer=tgt_norm,
             batch_size=4,
             shuffle=False,
-            conditioning="embedding",
+            particle_conditioning="embedding",
+            material_conditioning="embedding",
         )
         return sum(len(batch[0]) for batch in ds)
 

tests/test_dwarf.py

@@ -103,7 +103,7 @@ def test_warm_cache_writes_sidecar(tmp_path):
     assert loaded is not None
     assert loaded.vocab is not None
     assert loaded.event_index is not None
-    assert "valfrac=0.1_seed=0_cond=physical" in loaded.normalizers
+    assert "valfrac=0.1_seed=0_pcond=physical_mcond=physical" in loaded.normalizers
 
 
 def test_warm_cache_second_run_hits_cache(tmp_path):
@@ -167,5 +167,5 @@ def test_warm_cache_different_val_fraction_is_separate_entry(tmp_path):
     assert "fitting normalizer (streaming)" in result.output
     loaded = setup_cache.load(data, [data])
     assert loaded is not None
-    assert "valfrac=0.1_seed=0_cond=physical" in loaded.normalizers
-    assert "valfrac=0.3_seed=0_cond=physical" in loaded.normalizers
+    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

tests/test_flow.py

@@ -52,7 +52,7 @@ def test_sample_flow_shape():
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     sample, n_sec = sample_flow(_small_model(), cond_cont, cond_cat, steps=5)
     assert sample.shape == (B, 9)
-    assert n_sec.shape == (B,)
+    assert n_sec is not None and n_sec.shape == (B,)
 
 
 def test_ddpm_loss_nonneg():
@@ -69,4 +69,4 @@ def test_sample_ddim_shape():
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     sample, n_sec = sample_ddim(_small_model(), cond_cont, cond_cat, schedule, steps=5)
     assert sample.shape == (B, 9)
-    assert n_sec.shape == (B,)
+    assert n_sec is not None and n_sec.shape == (B,)

tests/test_migration_v02_v03.py

@@ -198,6 +198,56 @@ def test_migrate_legacy_model_config_shape():
     assert migrated["stage2_model"]["decoder"] == "one_shot"
 
 
+def test_migrate_legacy_model_config_nonzero_expert_dims_raises():
+    """docs/v0.3.0-followups.md item 8 regression: a v0.2 checkpoint's
+    model_config carrying a non-default expert_hidden_dim/expert_n_blocks
+    must fail loudly through this path too (§4.2) — not just
+    giant.config.migrate_config's parallel TOML-load path. Silently dropping
+    these keys (build_router's kwarg filtering) would resize the experts
+    instead of refusing."""
+    legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
+    legacy_cfg["router"] = {
+        "enabled": True,
+        "expert_hidden_dim": 128,
+        "expert_n_blocks": 0,
+    }
+    try:
+        net._migrate_legacy_model_config(legacy_cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "expert_hidden_dim" in str(e)
+
+
+def test_migrate_legacy_model_config_zero_expert_dims_dropped_silently():
+    legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
+    legacy_cfg["router"] = {
+        "enabled": True,
+        "expert_hidden_dim": 0,
+        "expert_n_blocks": 0,
+    }
+    migrated = net._migrate_legacy_model_config(legacy_cfg)
+    assert "expert_hidden_dim" not in migrated["stage1_model"]["router"]
+    assert "expert_n_blocks" not in migrated["stage1_model"]["router"]
+    assert "expert_hidden_dim" not in migrated["stage2_model"]["router"]
+    assert "expert_n_blocks" not in migrated["stage2_model"]["router"]
+
+
+def test_build_models_with_legacy_config_nonzero_expert_dims_raises():
+    """The same check must also fire through the actual caller,
+    build_models, not just the internal helper directly."""
+    legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
+    legacy_cfg["router"] = {
+        "enabled": True,
+        "expert_hidden_dim": 128,
+        "expert_n_blocks": 0,
+    }
+    try:
+        net.build_models(legacy_cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "expert_hidden_dim" in str(e)
+
+
 def test_build_models_accepts_new_nested_shape_unchanged():
     """A dict that already has a 'stage1_model' key (the new shape) is
     passed through build_models without going through the legacy migration

tests/test_network.py

@@ -1,5 +1,8 @@
+import copy
+
 import pytest
 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 (
     AttentionHistory,
@@ -9,6 +12,7 @@ from giant.model.network import (
     Stage1Model,
     Stage2Autoregressive,
     Stage2OneShot,
+    build_models,
     cat_col_layout,
     stage2_trunk_sec_dim,
     stage2_type_dim,
@@ -628,3 +632,52 @@ def test_stage2_autoregressive_history_step_matches_parallel_history_encoder():
 def test_stage2_autoregressive_init_history_cache_is_none_for_markov():
     model = _build_stage2_ar("physical", "wgan", history="markov")
     assert model.init_history_cache() is None
+
+
+# ── build_models: conditioning.share_stages ─────────────────────────────────
+
+
+def _minimal_model_config(share_stages: bool) -> dict:
+    cfg = copy.deepcopy(gconfig.DEFAULT_CONFIG)
+    cfg["conditioning"]["share_stages"] = share_stages
+    cfg["conditioning"]["particle"]["emb_dim"] = 4
+    cfg["conditioning"]["material"]["emb_dim"] = 4
+    cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1})
+    cfg["stage2_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "k_max": 3})
+    return {
+        "pdg_vocab": 3,
+        "mat_vocab": 2,
+        "conditioning": cfg["conditioning"],
+        "stage1_model": cfg["stage1_model"],
+        "stage2_model": cfg["stage2_model"],
+    }
+
+
+def test_build_models_share_stages_true_shares_condition_encoder_instance():
+    built = build_models(_minimal_model_config(share_stages=True))
+    stage1, stage2 = built["stage1"], built["stage2"]
+    assert stage1 is not None and stage2 is not None
+    assert stage1.cond_enc is stage2.cond_enc
+
+
+def test_build_models_share_stages_false_builds_independent_condition_encoders():
+    built = build_models(_minimal_model_config(share_stages=False))
+    stage1, stage2 = built["stage1"], built["stage2"]
+    assert stage1 is not None and stage2 is not None
+    assert stage1.cond_enc is not stage2.cond_enc
+
+
+def test_build_models_share_stages_true_shared_params_are_in_both_stage_parameter_lists():
+    """The shared encoder's parameters must actually appear in both stages'
+    own `.parameters()` — that's what makes each stage's independent
+    optimizer include (and update) them, which is the actual mechanism behind
+    "shared weights, forced common representation" (docs/v0.3.0-design.md
+    §3.1), not just object identity on `.cond_enc`."""
+    built = build_models(_minimal_model_config(share_stages=True))
+    stage1, stage2 = built["stage1"], built["stage2"]
+    assert stage1 is not None and stage2 is not None
+
+    shared_ids = {id(p) for p in stage1.cond_enc.parameters()}
+    assert shared_ids
+    assert shared_ids <= {id(p) for p in stage1.parameters()}
+    assert shared_ids <= {id(p) for p in stage2.parameters()}

tests/test_pipeline.py

@@ -7,6 +7,7 @@ import torch
 
 from giant import config as gconfig
 from giant.data import setup_cache
+from giant.data.transforms import Normalizer
 from giant.pipeline import run_train_job
 
 
@@ -235,6 +236,74 @@ def test_run_train_job_new_val_fraction_is_partial_hit(tmp_path, data, monkeypat
     assert "fitting normalizer (streaming)" in joined
 
 
+def test_run_train_job_custom_k_max_end_to_end(tmp_path, data):
+    """docs/v0.3.0-followups.md item 3 regression: stage2_model.k_max other
+    than the K_MAX module constant's default (15) must not produce a shape
+    mismatch between the data pipeline (loader.py/transforms.py padding) and
+    the model (network.py's trunks, sized from this same config value)."""
+    cfg = _tiny_cfg()
+    cfg["stage2_model"]["k_max"] = 3
+    _run(data, tmp_path / "out", cfg=cfg)
+    ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
+    assert ckpt["model_config"]["stage2_model"]["k_max"] == 3
+
+
+def test_run_train_job_mixed_particle_material_conditioning_end_to_end(tmp_path, data):
+    """docs/v0.3.0-followups.md item 4 regression: conditioning.particle.type
+    and conditioning.material.type are configured independently and may mix
+    freely (docs/v0.3.0-design.md §3.1) — e.g. particle "embedding" with
+    material "physical" — end-to-end through the real data pipeline, not
+    just accepted by validate_config."""
+    cfg = _tiny_cfg()
+    cfg["conditioning"]["particle"]["type"] = "embedding"
+    cfg["conditioning"]["material"]["type"] = "physical"
+    _run(data, tmp_path / "out", cfg=cfg)
+    ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
+    cond_cfg = ckpt["model_config"]["conditioning"]
+    assert cond_cfg["particle"]["type"] == "embedding"
+    assert cond_cfg["material"]["type"] == "physical"
+
+    cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
+    # Particle block ([COND_DIM_BASE:COND_DIM_BASE+PARTICLE_PHYS_DIM]) stays
+    # unfitted (mean=0/std=1) since "embedding" never computes real values
+    # for it; the material block is fit for real under "physical".
+    from giant.constants import COND_DIM_BASE, PARTICLE_PHYS_DIM
+
+    assert cond_norm.mean is not None and cond_norm.std is not None
+    np.testing.assert_allclose(
+        cond_norm.mean[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 0.0
+    )
+    np.testing.assert_allclose(
+        cond_norm.std[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 1.0
+    )
+    material_std = cond_norm.std[COND_DIM_BASE + PARTICLE_PHYS_DIM :]
+    assert np.all(material_std > 0) and not np.allclose(material_std, 1.0)
+
+
+def test_run_train_job_share_stages_end_to_end(tmp_path, data):
+    """docs/v0.3.0-followups.md item 5 regression: conditioning.share_stages
+    = true must actually train (not raise NotImplementedError), and the
+    resulting checkpoint's two stages must reload into a single shared
+    ConditionEncoder instance rather than two independent ones."""
+    from giant.model.network import build_models
+
+    cfg = _tiny_cfg()
+    cfg["conditioning"]["share_stages"] = True
+    _run(data, tmp_path / "out", cfg=cfg)
+    ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
+    assert ckpt["model_config"]["conditioning"]["share_stages"] is True
+
+    built = build_models(ckpt["model_config"])
+    stage1, stage2 = built["stage1"], built["stage2"]
+    assert stage1 is not None and stage2 is not None
+    assert stage1.cond_enc is stage2.cond_enc
+
+    stage1.load_state_dict(ckpt["model"])
+    stage2.load_state_dict(ckpt["sec_decoder"])
+    for p1, p2 in zip(stage1.cond_enc.parameters(), stage2.cond_enc.parameters()):
+        assert torch.equal(p1, p2)
+
+
 def test_run_train_job_matches_uncached_output(tmp_path, data):
     _run(data, tmp_path / "uncached", cache_setup=False)
     _run(data, tmp_path / "cached1", cache_setup=True)

tests/test_rollout.py

@@ -111,7 +111,8 @@ def _run(
         batch_size=128,
         max_tracks_per_event=max_tracks_per_event,
         escape_threshold=escape_threshold,
-        conditioning=conditioning,
+        particle_conditioning=conditioning,
+        material_conditioning=conditioning,
     )
 
 
@@ -172,7 +173,7 @@ def test_seed_frontier_embedding_mode_skips_unresolvable_pdg_lookup():
     frontier construction — mass/charge are simply zero-filled, unused."""
     seeds = _seeds(3)
     seeds["pdg"] = np.full(3, 999999999, dtype=np.int64)
-    fr, _counts = make_seed_frontier(**seeds, conditioning="embedding")
+    fr, _counts = make_seed_frontier(**seeds, particle_conditioning="embedding")
     np.testing.assert_array_equal(fr["mass"], 0.0)
     np.testing.assert_array_equal(fr["charge"], 0.0)
 
@@ -384,25 +385,43 @@ def _models_v3(
         noise_dim=8,
     )
     particle_type_cfg = {"target": target}
-    common = dict(
-        pdg_vocab=3,
-        mat_vocab=2,
-        particle_cfg=particle_cfg,
-        material_cfg=material_cfg,
-        hidden_dim=32,
-        n_res_blocks=2,
-        generator=generator2,
-        time_dim=16,
-        noise_dim=8,
-        k_max=k_max,
-        particle_type_cfg=particle_type_cfg,
-        build_n_sec_head=stage2_has_n_sec_head,
-    )
+    # 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
+    # keyword not itself part of that union (router, cond_enc, ...) look like
+    # a type mismatch to `ty` even though every actual value passed is fine.
     if decoder == "one_shot":
         sec_dim = stage2_trunk_sec_dim(particle_type_cfg, generator2, k_max, emb_dim)
-        s2 = Stage2OneShot(sec_dim=sec_dim, **common)
+        s2 = Stage2OneShot(
+            pdg_vocab=3,
+            mat_vocab=2,
+            particle_cfg=particle_cfg,
+            material_cfg=material_cfg,
+            hidden_dim=32,
+            n_res_blocks=2,
+            generator=generator2,
+            time_dim=16,
+            noise_dim=8,
+            k_max=k_max,
+            particle_type_cfg=particle_type_cfg,
+            build_n_sec_head=stage2_has_n_sec_head,
+            sec_dim=sec_dim,
+        )
     else:
-        s2 = Stage2Autoregressive(**common)
+        s2 = Stage2Autoregressive(
+            pdg_vocab=3,
+            mat_vocab=2,
+            particle_cfg=particle_cfg,
+            material_cfg=material_cfg,
+            hidden_dim=32,
+            n_res_blocks=2,
+            generator=generator2,
+            time_dim=16,
+            noise_dim=8,
+            k_max=k_max,
+            particle_type_cfg=particle_type_cfg,
+            build_n_sec_head=stage2_has_n_sec_head,
+        )
     return s1.eval(), s2.eval()
 
 
@@ -440,7 +459,8 @@ def _run_v3(
         batch_size=128,
         max_tracks_per_event=max_tracks_per_event,
         escape_threshold=escape_threshold,
-        conditioning=conditioning,
+        particle_conditioning=conditioning,
+        material_conditioning=conditioning,
         pdg_topn_map=pdg_topn_map,
         other_policy=other_policy,
         seed=seed,
@@ -535,6 +555,88 @@ def test_rollout_onehot_target_missing_topn_map_raises(fake_material_props):
         _run_v3(s1, s2, pdg_topn_map=None)
 
 
+# --- conditioning.{particle,material}.type = "onehot" (docs/v0.3.0-followups.md
+# item 7) — a separate axis from stage2_model.particle_type.target above: this
+# is what feeds cond_cat's extra top-N columns for ConditionEncoder's own
+# "onehot" mode, not the secondary-species decode. ---------------------------
+
+COND_PDG_TOPN_MAP = TopNMap(class_map=dict(PDG_MAP), other_members={})
+COND_MAT_TOPN_MAP = TopNMap(class_map={"G4_AIR": 0, "G4_PbWO4": 1}, other_members={})
+
+
+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}
+    s1 = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+    )
+    s2 = Stage2OneShot(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+        sec_dim=stage2_trunk_sec_dim({"target": "physical"}, "flow", K_MAX, 3),
+        generator="flow",
+        time_dim=16,
+    )
+    return s1.eval(), s2.eval()
+
+
+def _run_onehot_conditioning(
+    pdg_topn_map=COND_PDG_TOPN_MAP, mat_topn_map=COND_MAT_TOPN_MAP
+):
+    s1, s2 = _onehot_conditioning_models()
+    cond, tgt, sec_phys = _norms()
+    return rollout(
+        s1,
+        s2,
+        _oracle(),
+        _seeds(),
+        cond,
+        tgt,
+        sec_phys,
+        PDG_MAP,
+        MAT_MAP,
+        energy_cutoff=1.0,
+        max_steps=30,
+        steps=4,
+        batch_size=128,
+        max_tracks_per_event=300,
+        escape_threshold=1e9,
+        particle_conditioning="onehot",
+        material_conditioning="onehot",
+        pdg_topn_map=pdg_topn_map,
+        mat_topn_map=mat_topn_map,
+    )
+
+
+def test_rollout_conditioning_onehot_end_to_end(fake_material_props):
+    rec = _run_onehot_conditioning()
+    assert len(rec["event_id"]) > 0
+    assert set(rec["event_id"].tolist()) == set(range(6))
+
+
+def test_rollout_conditioning_onehot_particle_missing_topn_map_raises(
+    fake_material_props,
+):
+    with pytest.raises(RuntimeError, match="pdg_topn_map"):
+        _run_onehot_conditioning(pdg_topn_map=None)
+
+
+def test_rollout_conditioning_onehot_material_missing_topn_map_raises(
+    fake_material_props,
+):
+    with pytest.raises(RuntimeError, match="mat_topn_map"):
+        _run_onehot_conditioning(mat_topn_map=None)
+
+
 @pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"])
 def test_rollout_embedding_target_end_to_end(decoder):
     """particle_type.target="embedding" L1-snaps to the nearest row of the

tests/test_sample.py

@@ -150,7 +150,7 @@ def test_sample_flow_returns_n_sec_for_legacy_stage1():
     )
     cond_cont, cond_cat = _cond(5)
     _, n_sec = sample_flow(model, cond_cont, cond_cat, steps=2)
-    assert n_sec.shape == (5,)
+    assert n_sec is not None and n_sec.shape == (5,)
 
 
 # ── Stage2OneShot: non-"physical" particle_type.target ──────────────────────

tests/test_transforms.py

@@ -426,7 +426,13 @@ def test_build_features_embedding_mode_zero_fills_physical_columns():
     data = _minimal_step_data(3)
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
 
-    cond_cont, *_ = build_features(data, pdg_map, mat_map, conditioning="embedding")
+    cond_cont, *_ = build_features(
+        data,
+        pdg_map,
+        mat_map,
+        particle_conditioning="embedding",
+        material_conditioning="embedding",
+    )
 
     assert cond_cont.shape[1] == COND_DIM
     np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE:], 0.0)
@@ -438,7 +444,13 @@ def test_build_features_physical_mode_shape_and_values(fake_material_props):
     data = _minimal_step_data(3)
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
 
-    cond_cont, *_ = build_features(data, pdg_map, mat_map, conditioning="physical")
+    cond_cont, *_ = build_features(
+        data,
+        pdg_map,
+        mat_map,
+        particle_conditioning="physical",
+        material_conditioning="physical",
+    )
 
     assert cond_cont.shape[1] == COND_DIM
     mass, charge = particle_mass_charge(11)
@@ -461,7 +473,13 @@ def test_build_features_physical_mode_unfilled_material_raises():
     pdg_map, mat_map = {11: 0}, {"G4_LYSO": 0}
 
     with pytest.raises(MaterialPropertiesNotFilledError):
-        build_features(data, pdg_map, mat_map, conditioning="physical")
+        build_features(
+            data,
+            pdg_map,
+            mat_map,
+            particle_conditioning="physical",
+            material_conditioning="physical",
+        )
 
 
 def test_build_cond_features_mass_charge_override(fake_material_props):
@@ -474,7 +492,13 @@ def test_build_cond_features_mass_charge_override(fake_material_props):
     data["charge"] = np.array([2.0, -2.0], dtype=np.float32)
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
 
-    cond_cont, _ = build_cond_features(data, pdg_map, mat_map, conditioning="physical")
+    cond_cont, _ = build_cond_features(
+        data,
+        pdg_map,
+        mat_map,
+        particle_conditioning="physical",
+        material_conditioning="physical",
+    )
 
     np.testing.assert_allclose(
         cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0]))
@@ -494,7 +518,12 @@ def test_build_cond_features_pads_legacy_normalizer_in_embedding_mode():
     legacy_norm.std = np.ones(COND_DIM_BASE, dtype=np.float32)
 
     cond_cont, _ = build_cond_features(
-        data, pdg_map, mat_map, cond_normalizer=legacy_norm, conditioning="embedding"
+        data,
+        pdg_map,
+        mat_map,
+        cond_normalizer=legacy_norm,
+        particle_conditioning="embedding",
+        material_conditioning="embedding",
     )
 
     assert cond_cont.shape[-1] == COND_DIM
@@ -521,7 +550,8 @@ def test_build_cond_features_rejects_legacy_normalizer_in_physical_mode(
             pdg_map,
             mat_map,
             cond_normalizer=legacy_norm,
-            conditioning="physical",
+            particle_conditioning="physical",
+            material_conditioning="physical",
         )
 
 
@@ -610,13 +640,23 @@ def test_build_cond_features_physical_mode_tolerates_out_of_vocab_pdg_and_materi
     }
 
     cond_cont, cond_cat = build_cond_features(
-        data, pdg_map, mat_map, conditioning="physical"
+        data,
+        pdg_map,
+        mat_map,
+        particle_conditioning="physical",
+        material_conditioning="physical",
     )
     assert cond_cont.shape[-1] == COND_DIM
     np.testing.assert_array_equal(cond_cat, [[0, 0]])  # dummy indices, no raise
 
     with pytest.raises(KeyError):
-        build_cond_features(data, pdg_map, mat_map, conditioning="embedding")
+        build_cond_features(
+            data,
+            pdg_map,
+            mat_map,
+            particle_conditioning="embedding",
+            material_conditioning="embedding",
+        )
 
 
 # ── _WelfordAccumulator ──────────────────────────────────────────────────────

tests/test_validate.py

@@ -1,15 +1,18 @@
 import numpy as np
 import torch
 
-from giant.constants import COND_DIM, K_MAX, SEC_SLOT_DIM, X_DIM
-from giant.model.network import Stage1Model, Stage2OneShot
+from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM
+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}
+_K_MAX = 5
 
 
-def _tiny_models():
+def _tiny_models(particle_type_cfg: dict | None = None):
+    """A fresh v0.3.0 pair: Stage1Model owns no n_sec_head (decision 1), so
+    n_sec always comes from Stage2OneShot."""
     s1 = Stage1Model(
         pdg_vocab=3,
         mat_vocab=2,
@@ -17,7 +20,13 @@ def _tiny_models():
         material_cfg=_MATERIAL_CFG,
         hidden_dim=16,
         n_res_blocks=1,
-        n_sec_head_k_max=K_MAX,
+    )
+    target = (particle_type_cfg or {}).get("target", "physical")
+    sec_dim = stage2_trunk_sec_dim(
+        particle_type_cfg or {"target": "physical"},
+        "flow",
+        _K_MAX,
+        int(_PARTICLE_CFG["emb_dim"]),
     )
     s2 = Stage2OneShot(
         pdg_vocab=3,
@@ -28,23 +37,29 @@ def _tiny_models():
         n_res_blocks=1,
         generator="flow",
         time_dim=16,
+        k_max=_K_MAX,
+        sec_dim=sec_dim,
+        particle_type_cfg=particle_type_cfg,
     )
+    assert s2.particle_type_cfg.get("target", "physical") == target
     return s1.eval(), s2.eval()
 
 
-def _zero_secondaries_loader(B=4, n_batches=2):
-    """A val_loader whose every batch has n_sec=0 (real side) — matches the
-    (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx) tuple shape
-    StreamingStepsDataset yields."""
+def _loader(B: int = 4, n_batches: int = 2, n_sec_value: int = 0, n_classes: int = 8):
+    """A val_loader matching StreamingStepsDataset's 7-tuple batch shape:
+    (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, sec_type_idx)."""
     batches = []
     for _ in range(n_batches):
         cond_cont = torch.randn(B, COND_DIM)
         cond_cat = torch.zeros(B, 2, dtype=torch.long)
         x1 = torch.randn(B, X_DIM)
-        n_sec = torch.zeros(B, dtype=torch.long)
-        sec_cont = torch.zeros(B, K_MAX, SEC_SLOT_DIM)
+        n_sec = torch.full((B,), n_sec_value, dtype=torch.long)
+        sec_cont = torch.randn(B, _K_MAX, SEC_SLOT_DIM)
         proc_idx = torch.zeros(B, dtype=torch.long)
-        batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx))
+        sec_type_idx = torch.randint(0, n_classes, (B, _K_MAX), dtype=torch.long)
+        batches.append(
+            (cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx)
+        )
     return batches
 
 
@@ -52,22 +67,56 @@ def test_validate_marginals_all_zero_secondaries_returns_nan_phys_kl(monkeypatch
     """If n_sec_pred collapses to 0 across the whole validated set (realistic
     during early/unstable training), phys_kl must degrade to NaN instead of
     crashing on the empty-array .min()/.max() reduction inside
-    _histogram_kl -- a regression the old species/bincount code this
-    replaced explicitly guarded against."""
+    _histogram_kl."""
     s1, s2 = _tiny_models()
-    loader = _zero_secondaries_loader()
+    loader = _loader(n_sec_value=0)
 
-    # Force the Stage-1 n_sec head's prediction to 0 for every sample too, so
-    # the generated side's valid-slot mask is also empty (real side is
-    # already all n_sec=0 by construction of the fake loader above).
-    def _fake_sample_flow(model, cond_cont, cond_cat, **kw):
-        B = cond_cont.size(0)
-        return torch.randn(B, X_DIM), torch.zeros(B, dtype=torch.long)
+    def _fake_resolve_n_sec(
+        stage1_model, sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred
+    ):
+        return torch.zeros(cond_cont.size(0), dtype=torch.long)
 
-    monkeypatch.setattr("giant.validate.sample_flow", _fake_sample_flow)
+    monkeypatch.setattr("giant.validate.resolve_n_sec", _fake_resolve_n_sec)
 
-    result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=2)
+    result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=2, steps=2)
 
     assert np.asarray(result["phys_real"]).shape == (0, 2)
     assert np.asarray(result["phys_generated"]).shape == (0, 2)
     assert np.isnan(np.asarray(result["phys_kl"])).all()
+
+
+def test_validate_marginals_physical_target_shapes():
+    s1, s2 = _tiny_models()
+    loader = _loader(n_sec_value=2)
+
+    result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=1, steps=2)
+
+    assert np.asarray(result["real"]).shape == (4, X_DIM)
+    assert np.asarray(result["generated"]).shape == (4, X_DIM)
+    assert np.asarray(result["kl_divergence"]).shape == (X_DIM,)
+    assert "phys_real" in result and "phys_generated" in result and "phys_kl" in result
+    assert "type_class_real" not in result
+
+
+def test_validate_marginals_onehot_type_class_marginal():
+    particle_type_cfg = {"target": "onehot"}
+    s1, s2 = _tiny_models(particle_type_cfg)
+    loader = _loader(n_sec_value=2, n_classes=s2.type_dim)
+
+    result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=1, steps=2)
+
+    assert "phys_real" not in result
+    # Real/generated valid-slot counts need not agree (real: ground-truth
+    # n_sec=2 always; generated: the untrained n_sec_head's own prediction).
+    assert np.asarray(result["type_class_real"]).ndim == 1
+    assert np.asarray(result["type_class_gen"]).ndim == 1
+    assert np.asarray(result["type_class_real"]).shape[0] > 0
+
+
+def test_validate_marginals_without_sec_decoder_returns_stage1_only():
+    s1, _ = _tiny_models()
+    loader = _loader(n_sec_value=0)
+
+    result = validate_marginals(s1, loader, n_batches=1, steps=2)
+
+    assert set(result) == {"real", "generated", "kl_divergence"}

tests/test_wgan.py

@@ -120,7 +120,7 @@ def test_sample_wgan_shape():
     cond_cont, cond_cat = _cond(B)
     sample, n_sec = sample_wgan(model, cond_cont, cond_cat)
     assert sample.shape == (B, X_DIM)
-    assert n_sec.shape == (B,)
+    assert n_sec is not None and n_sec.shape == (B,)
 
 
 # --- Stage-2 generator/critic ---