v0.3.0 step 6: sample.py/rollout.py AR generation + class->PDG decode

- giant/sample.py: fix every sampler's call convention against
  Stage1Model/Stage2OneShot's actual forward signatures (was still
  calling model(x, t, cond_cont, cond_cat) positionally); add
  sample_secondaries_ar (free-running AR loop, unsnapped history feature)
  and sample_stage1/sample_stage2/resolve_n_sec dispatch helpers that read
  each stage's generator_kind/decoder off the model instance itself.
- giant/particles.py: decode_topn_class (argmax + other_policy) and
  decode_embedding_nearest (L1-snap + distance) turn a secondary's
  "onehot"/"embedding" type prediction into a concrete PDG.
- giant/rollout.py: decode_secondary_identity routes all three
  particle_type.target values to real mass/charge; per-stage generator
  dispatch (drops the single shared `mode` string, adds ddpm support);
  L1DistCollector accumulates the §11.3 embedding-distance diagnostic.
- giant/cli.py: drop the onehot/embedding-target rejection gate (narrowed
  to the still-unimplemented conditioning.particle/material.type=onehot
  axis); fix the dead model_cfg.get("mode") bug in predict/rollout.
- giant/analysis/: new type_embedding_l1_distance PlotSpec, wired through
  the rollout YAML sidecar (no live-model call needed, unlike
  router_gating -- the histogram is already pre-aggregated at rollout
  time).
- Un-xfail every test that was blocked on this step (test_rollout.py,
  test_flow.py, test_wgan.py, test_phase2.py, test_router.py,
  test_validate.py); add test_sample.py, test_type_embedding_distance.py.

Known follow-up: giant/validate.py still unpacks the training val-batch
as a stale 6-tuple and doesn't use the new per-stage dispatch, so
marginal validation during training degrades gracefully with a warning
rather than working -- not in this step's scope.

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

giant/analysis/catalog.py

@@ -58,6 +58,7 @@ from giant.analysis.router_gating import (
     compute_router_share_by_process,
 )
 from giant.analysis.sources import Side, open_side, physical_steps, secondaries
+from giant.analysis.type_embedding_distance import compute_type_embedding_l1_distance
 from giant.analysis.variables import RANGED_VARS, cos_scatter_expr
 
 
@@ -71,6 +72,11 @@ class Bundle:
     r_phys: pl.LazyFrame  # rollout, physical steps only
     t_phys: pl.LazyFrame  # reference, physical steps only
     checkpoint: str | None = None  # from the rollout YAML; router_gating only
+    # §11.3 diagnostic pre-aggregated at rollout time (giant.rollout.
+    # L1DistCollector.summary()) — from the rollout YAML, type_embedding_l1_distance
+    # only. Unlike checkpoint/router_gating, this needs no live model: it's
+    # already a finished histogram, just passed through.
+    type_embedding_l1_dist: dict | None = None
 
     @classmethod
     def open(
@@ -80,6 +86,7 @@ class Bundle:
         ctx: Context,
         checkpoint=None,
         chunk: tuple[int, int] | None = None,
+        type_embedding_l1_dist: dict | None = None,
     ) -> "Bundle":
         """Open both sides, optionally restricted to one event-disjoint chunk.
 
@@ -103,6 +110,7 @@ class Bundle:
             r_phys=physical_steps(r_all, Side.rollout),
             t_phys=physical_steps(t_all, Side.reference),
             checkpoint=checkpoint,
+            type_embedding_l1_dist=type_embedding_l1_dist,
         )
 
 
@@ -658,6 +666,9 @@ _router_share_pdg_partial, _router_share_pdg_finalize = _unchunkable(
 _router_share_process_partial, _router_share_process_finalize = _unchunkable(
     lambda b: compute_router_share_by_process(b.checkpoint, b.t_phys)
 )
+_type_embedding_l1_distance_partial, _type_embedding_l1_distance_finalize = (
+    _unchunkable(lambda b: compute_type_embedding_l1_distance(b.type_embedding_l1_dist))
+)
 
 
 # ---------------------------------------------------------------------------
@@ -831,6 +842,13 @@ def build_catalog() -> list[PlotSpec]:
             finalize=_router_share_process_finalize,
             chunkable=False,
         ),
+        PlotSpec(
+            "type_embedding_l1_distance",
+            "model",
+            compute_partial=_type_embedding_l1_distance_partial,
+            finalize=_type_embedding_l1_distance_finalize,
+            chunkable=False,
+        ),
     ]
     return specs
 

giant/analysis/condor.py

@@ -83,6 +83,12 @@ _PLOT_META_KEYS = (
     "best_val_loss",
     "training_config",
     "training_meta",
+    # §11.3 diagnostic — only present when giant rollout ran under
+    # stage2_model.particle_type.target="embedding" (see giant/cli.py's
+    # rollout command and giant.rollout.L1DistCollector); absent otherwise,
+    # which the type_embedding_l1_distance PlotSpec (catalog.py) reads as
+    # "not applicable to this checkpoint".
+    "type_embedding_l1_dist",
 )
 
 
@@ -247,6 +253,7 @@ def compute_reduced(
     checkpoint: str | None = None,
     chunk_index: int = 0,
     n_chunks: int = 1,
+    type_embedding_l1_dist: dict | None = None,
 ) -> Path:
     """Core: run one (plot, chunk)'s partial reduction against explicit paths.
 
@@ -265,7 +272,12 @@ def compute_reduced(
             f"n_chunks={effective_n} (chunkable={spec.chunkable})"
         )
     bundle = Bundle.open(
-        rollout, reference, ctx, checkpoint=checkpoint, chunk=(chunk_index, effective_n)
+        rollout,
+        reference,
+        ctx,
+        checkpoint=checkpoint,
+        chunk=(chunk_index, effective_n),
+        type_embedding_l1_dist=type_embedding_l1_dist,
     )
     partial = Partial(
         id=spec_id,
@@ -291,6 +303,7 @@ def compute_one(spec_id: str, run_dir: str | Path, chunk_index: int = 0) -> Path
         checkpoint=meta.plot_meta.get("checkpoint"),
         chunk_index=chunk_index,
         n_chunks=meta.n_chunks,
+        type_embedding_l1_dist=meta.plot_meta.get("type_embedding_l1_dist"),
     )
 
 

giant/analysis/render.py

@@ -144,6 +144,8 @@ def _render_single(r: Reduced, params: dict):
     )
     if r.payload.get("log_y"):
         ax.set_yscale("log")
+    if r.payload.get("log_x"):
+        ax.set_xscale("log")
     ax.set_xlabel(r.xlabel)
     ax.set_ylabel("density")
     ps.style_legend(ax, title="source")

giant/analysis/type_embedding_distance.py

@@ -0,0 +1,70 @@
+"""Secondary-type embedding-distance diagnostic (docs/v0.3.0-design.md §11.3).
+
+Unlike every other diagnostic in this package, the data isn't derivable from
+a rollout/reference parquet at all — it's the L1 distance between each
+emitted secondary's *raw* predicted embedding vector (under
+`stage2_model.particle_type.target = "embedding"`) and the nearest row of the
+conditioning's embedding table it snapped to, which only exists transiently
+inside `giant rollout` (`giant.rollout.decode_secondary_identity`), never
+written to a column. So it's accumulated once, at rollout time
+(`giant.rollout.L1DistCollector`), and stashed as a pre-finished histogram
+summary in the rollout YAML sidecar (`type_embedding_l1_dist`) — this module
+just turns that summary into a `Reduced`, no parquet scan involved (a
+`chunkable=False` spec, like `router_gating`, but even cheaper: no live model
+call either).
+
+A heavy right tail means the decoder is emitting vectors off the embedding
+manifold — the direct analogue of the species-collapse symptom the v0.3.0
+redesign exists to fix.
+"""
+
+from __future__ import annotations
+
+from giant.analysis.reduced import Reduced
+
+_NOTE_NOT_APPLICABLE = (
+    "not applicable: this rollout's checkpoint doesn't use "
+    "stage2_model.particle_type.target='embedding' (or generated no "
+    "secondaries), so giant rollout recorded no type_embedding_l1_dist "
+    "diagnostic in its YAML sidecar"
+)
+
+
+def compute_type_embedding_l1_distance(l1_dist: dict | None) -> Reduced:
+    """`Reduced` for the type-embedding-distance figure, or an explanatory
+    note if this checkpoint never populated the diagnostic.
+
+    `l1_dist`: `giant.rollout.L1DistCollector.summary()`'s dict, as recorded
+    in the rollout YAML's `type_embedding_l1_dist` key (`Bundle.
+    type_embedding_l1_dist`) — `{"n", "mean", "std", "min", "max",
+    "hist_edges", "hist_counts"}`.
+    """
+    if l1_dist is None:
+        return Reduced(
+            id="type_embedding_l1_distance",
+            family="model",
+            kind="unavailable",
+            title="Secondary-type embedding L1 distance",
+            xlabel="n/a",
+            payload={"note": _NOTE_NOT_APPLICABLE},
+        )
+
+    return Reduced(
+        id="type_embedding_l1_distance",
+        family="model",
+        kind="single_hist",
+        title="Secondary-type embedding L1 distance (predicted vector -> nearest PDG row)",
+        xlabel="L1 distance",
+        payload={
+            "edges": l1_dist["hist_edges"],
+            "rollout": l1_dist["hist_counts"],
+            "log_y": True,
+            "log_x": True,
+            "note": (
+                f"n={l1_dist['n']:,}  mean={l1_dist['mean']:.4g}  "
+                f"std={l1_dist['std']:.4g}  min={l1_dist['min']:.4g}  "
+                f"max={l1_dist['max']:.4g}; rollout only, no reference "
+                "concept for a raw pre-decode vector"
+            ),
+        },
+    )

giant/cli.py

@@ -34,24 +34,22 @@ from giant.data.loader import (
 from giant.data.transforms import (
     build_features,
     build_cond_features,
-    decode_secondaries,
     energy_simplex_decode,
     inv_local_frame_rotation,
     inv_log_transform,
     reconstruct_post_pos,
     Normalizer,
 )
+from giant.data.setup_cache import topnmap_from_json
 from giant.geometry import GeometryOracle
 from giant.model.network import build_models
-from giant.particles import nearest_known_pdg
 from giant.pipeline import run_train_job
-from giant.rollout import rollout as run_rollout
-from giant.sample import (
-    sample_flow,
-    sample_secondaries,
-    sample_wgan,
-    sample_secondaries_wgan,
+from giant.rollout import (
+    L1DistCollector,
+    decode_secondary_identity,
+    rollout as run_rollout,
 )
+from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
 
 app = typer.Typer(no_args_is_help=True)
 
@@ -86,41 +84,66 @@ def _conditioning_str(model_cfg: dict, default: str = "embedding") -> str:
     return raw
 
 
-def _check_v030_onehot_support(model_cfg: dict, command: str) -> None:
-    """`giant predict`/`giant rollout` don't yet support conditioning
-    `"onehot"` mode or `stage2_model.particle_type.target` in `("onehot",
-    "embedding")` — full decode (class index -> concrete PDG, `other_policy`
-    sampling) is v0.3.0 step 6 (docs/v0.3.0-design.md), which also lands the
-    autoregressive decoder these targets are meant to pair with. Without
-    this guard, predict/rollout would crash later on a `cond_cat`/`sec_dim`
-    shape mismatch (onehot) or silently read a meaningless raw (log_mass,
-    charge) pair (embedding) instead of failing clearly. A v0.2 (flat)
-    `model_config` never has these, so this is a no-op there.
+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")
-    particle_type_target = (
-        model_cfg.get("stage2_model", {}).get("particle_type", {}).get("target")
-    )
-    if (
-        particle_type == "onehot"
-        or material_type == "onehot"
-        or particle_type_target in ("onehot", "embedding")
-    ):
+    if particle_type == "onehot" or material_type == "onehot":
         typer.echo(
-            f"error: giant {command} does not yet support conditioning "
-            "onehot mode or stage2_model.particle_type.target in "
-            "('onehot', 'embedding') — full decode (class index -> concrete "
-            "PDG) lands in v0.3.0 step 6 alongside the autoregressive "
-            "decoder these targets are meant to pair with.",
+            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,
         )
         raise typer.Exit(1)
 
 
+def _stage_cfg(model_cfg: dict, stage: str) -> dict:
+    """`model_cfg[f"{stage}_model"]` for a new-format model_config, `{}` for
+    a v0.2 flat one (whose ddpm schedule always used `CosineSchedule`'s own
+    default `T=1000` — never a config key — and which never had
+    `particle_type` at all, so `{}` is the correct fallback for both
+    `_ddpm_steps`/`_particle_type_other_policy` below)."""
+    val = model_cfg.get(f"{stage}_model")
+    return val if isinstance(val, dict) else {}
+
+
+def _ddpm_steps(model_cfg: dict, stage: str) -> int:
+    return _stage_cfg(model_cfg, stage).get("ddpm", {}).get("n_steps", 1000)
+
+
+def _particle_type_other_policy(model_cfg: dict) -> str:
+    return (
+        _stage_cfg(model_cfg, "stage2")
+        .get("particle_type", {})
+        .get("other_policy", "sample")
+    )
+
+
+def _load_pdg_topn_map(ckpt: dict):
+    """`ckpt["pdg_topn_map"]` as a `giant.data.loader.TopNMap`, or `None` if
+    this checkpoint's conditioning/particle_type never needed one (see
+    `giant.pipeline.run_setup_stage`, which only populates it when
+    `conditioning.particle.type` or `stage2_model.particle_type.target` is
+    `"onehot"`)."""
+    raw = ckpt.get("pdg_topn_map")
+    return topnmap_from_json(raw, axis="pdg") if raw is not None else None
+
+
 def _batch_size_estimate_dims(
     model_cfg: dict, training: bool, stage: str = "stage1"
 ) -> tuple[int, int]:
@@ -965,7 +988,10 @@ def predict(
         raise typer.Exit(1)
 
     model_cfg = ckpt["model_config"]
-    _check_v030_onehot_support(model_cfg, "predict")
+    _check_conditioning_onehot_support(model_cfg, "predict")
+    pdg_topn_map = _load_pdg_topn_map(ckpt)
+    other_policy = _particle_type_other_policy(model_cfg)
+    stage1_ddpm_steps = _ddpm_steps(model_cfg, "stage1")
 
     if batch_size_auto:
         est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(
@@ -1046,23 +1072,21 @@ def predict(
 
         cc = torch.from_numpy(cond_cont).float().to(_device)
         ck = torch.from_numpy(cond_cat).long().to(_device)
-        if model_cfg.get("mode") == "wgan":
-            stage1_norm, n_sec_pred = sample_wgan(model, cc, ck)
-        else:
-            stage1_norm, n_sec_pred = sample_flow(model, cc, ck, steps=steps)
+        stage1_norm, n_sec_pred = sample_stage1(
+            model, cc, ck, steps=steps, ddpm_steps=stage1_ddpm_steps
+        )
 
         if coord == Coord.global_:
-            if model_cfg.get("mode") == "wgan":
-                sec_cont, sec_phys, _sec_valid_pred = sample_secondaries_wgan(
-                    sec_decoder, cc, ck, stage1_norm, n_sec_pred
-                )
-            else:
-                sec_cont, sec_phys, _sec_valid_pred = sample_secondaries(
-                    sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps
-                )
-            sec_full_np = torch.cat([sec_cont, sec_phys], dim=-1).cpu().numpy()
+            # A fresh v0.3.0 Stage1Model owns no n_sec_head (decision 1) —
+            # sample_stage1 returns n_sec_pred=None then, so ask stage 2.
+            n_sec_pred = resolve_n_sec(
+                model, sec_decoder, cc, ck, stage1_norm, n_sec_pred
+            )
+            sec_cont, sec_type, _sec_valid_pred = sample_stage2(
+                sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps
+            )
+            n_sec_pred_np = n_sec_pred.cpu().numpy()
 
-        n_sec_pred_np = n_sec_pred.cpu().numpy()
         pred = stage1_norm.cpu().numpy()  # normalised
 
         # Inverse-normalise → local frame, log-scaled scalars
@@ -1119,19 +1143,26 @@ def predict(
                 piece["pre_pos"], piece["pre_dir"], step_length, travel_dir_local
             )
 
-            sec_E, sec_dir_world, sec_mass, sec_charge, _sec_valid = decode_secondaries(
-                sec_full_np,
-                n_sec_pred_np,
-                e_sec_pred,
-                piece["pre_dir"],
-                sec_phys_normalizer=sec_phys_norm,
+            # particle_type.target="physical": sec_pdg_code is a reporting-
+            # only nearest-known-PDG label (never fed back into the model —
+            # "no snapping at inference"). "onehot"/"embedding": PDG
+            # resolution IS the secondary's identity — see
+            # decode_secondary_identity's docstring.
+            sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, _l1_dist = (
+                decode_secondary_identity(
+                    sec_decoder,
+                    sec_cont,
+                    sec_type,
+                    n_sec_pred_np,
+                    e_sec_pred,
+                    piece["pre_dir"],
+                    sec_phys_norm,
+                    pdg_map,
+                    pdg_topn_map,
+                    other_policy,
+                    None,
+                )
             )
-            # Reporting-only nearest-known-PDG label (never fed back into the
-            # model) for the sec_pdg_list output column — see
-            # giant/particles.py and the "no snapping at inference" design.
-            sec_pdg_code = nearest_known_pdg(
-                sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()
-            ).reshape(sec_mass.shape)
             sec_pdg_list = [
                 sec_pdg_code[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)
             ]
@@ -1378,8 +1409,12 @@ def rollout(
     training_cfg = gconfig.load_checkpoint_config(checkpoint)
 
     model_cfg = ckpt["model_config"]
-    _check_v030_onehot_support(model_cfg, "rollout")
+    _check_conditioning_onehot_support(model_cfg, "rollout")
     conditioning = _conditioning_str(model_cfg)
+    pdg_topn_map = _load_pdg_topn_map(ckpt)
+    other_policy = _particle_type_other_policy(model_cfg)
+    stage1_ddpm_steps = _ddpm_steps(model_cfg, "stage1")
+    stage2_ddpm_steps = _ddpm_steps(model_cfg, "stage2")
     pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
     mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
     cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
@@ -1432,6 +1467,10 @@ def rollout(
             writer = pq.ParquetWriter(out, table.schema)
         writer.write_table(table)
 
+    # Only meaningful under particle_type.target="embedding" (§11.3) — a
+    # no-op collector otherwise, cheaper than branching the call itself.
+    l1_dist_collector = L1DistCollector()
+
     summary = run_rollout(
         model,
         sec_decoder,
@@ -1451,11 +1490,18 @@ def rollout(
         escape_threshold=escape_threshold,
         on_chunk=_write_chunk,
         conditioning=conditioning,
-        mode=model_cfg.get("mode", "flow"),
+        pdg_topn_map=pdg_topn_map,
+        other_policy=other_policy,
+        seed=seed,
+        stage1_ddpm_steps=stage1_ddpm_steps,
+        stage2_ddpm_steps=stage2_ddpm_steps,
+        l1_dist_collector=l1_dist_collector,
     )
     if writer is not None:
         writer.close()
 
+    l1_summary = l1_dist_collector.summary()
+
     ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path)
     ref = yaml.safe_load(ref_path.read_text())
     ref.update(
@@ -1475,6 +1521,13 @@ def rollout(
             "rollout_seed": seed,
             "n_rows": summary["n_rows"],
             "termination_reason_counts": summary["termination_reason_counts"],
+            # §11.3 diagnostic — only present under
+            # stage2_model.particle_type.target="embedding"; omitted (not
+            # written as null) otherwise, so giant.analysis can tell "not
+            # applicable to this checkpoint" apart from "collector empty".
+            **(
+                {"type_embedding_l1_dist": l1_summary} if l1_summary is not None else {}
+            ),
             # Full architecture spec baked into the checkpoint — includes the
             # entire router sub-dict, not just a hand-picked subset, so any
             # model knob (router type/n_experts, noise_dim, vocab sizes, ...)

giant/data/transforms.py

@@ -587,42 +587,37 @@ def encode_secondary_type_idx(
     return np.where(sec_valid, idx, 0).astype(np.int64)
 
 
-def decode_secondaries(
+def decode_secondary_cont(
     sec_cont: np.ndarray,
     n_sec: np.ndarray,
     e_sec: np.ndarray,
     pre_dir: np.ndarray,
-    sec_phys_normalizer: "Normalizer | None" = None,
-) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
-    """Inverse of encode_secondaries: continuous targets → physical secondary attrs.
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Continuous-only half of `decode_secondaries`'s inverse: the
+    stick-breaking energy split and local->world direction — generator/
+    `particle_type.target`-independent, since every target (`"physical"`,
+    `"onehot"`, `"embedding"`) shares the same `CONT_SLOT_DIM`-wide
+    (stick_logit, dir) prefix (docs/v0.3.0-design.md §6.1) and differs only
+    in what follows it. `decode_secondaries` (target="physical") is the
+    original all-in-one form built on top of this; `target` in `("onehot",
+    "embedding")` decodes their type slice separately via
+    `giant.particles.decode_topn_class`/`decode_embedding_nearest` and calls
+    this directly instead — see `giant/rollout.py`.
 
-    sec_cont: (N, K_MAX, 6) — [stick_logit, local_dir_x, local_dir_y,
-        local_dir_z, log_mass, charge] (log_mass/charge normalised iff
-        `sec_phys_normalizer` was applied when this was produced — e.g. a
-        raw model prediction; pass the same normalizer here to invert it)
+    sec_cont: (N, K, >=CONT_SLOT_DIM) — only columns `[:, :, :CONT_SLOT_DIM]`
+        (stick_logit, local dir) are read; a caller may pass its full
+        per-slot tensor (continuous + type) unsliced.
     n_sec: (N,) integer secondary counts
     e_sec: (N,) total secondary energy budget [MeV]
     pre_dir: (N, 3) pre-step world-frame direction
 
-    Returns (sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid) each
-    shape (N, K_MAX). The valid slots' energies (`sec_E[sec_valid]`, per row)
-    always sum to exactly `e_sec` — see the rescaling below. mass/charge are
-    the model's raw predicted physical identity for each secondary, used
-    as-is (no snapping to a discrete PDG code) — see giant/particles.py for
-    the separate, reporting-only nearest-PDG lookup callers may apply on top
-    of this for display/bookkeeping purposes.
+    Returns (sec_E, sec_dir_world, sec_valid), shapes (N, K), (N, K, 3),
+    (N, K). The valid slots' energies (`sec_E[sec_valid]`, per row) always
+    sum to exactly `e_sec` — see the rescaling below.
     """
-    if sec_phys_normalizer is not None:
-        N_, K_, _ = sec_cont.shape
-        phys = sec_phys_normalizer.inverse_transform(sec_cont[:, :, 4:6].reshape(-1, 2))
-        sec_cont = sec_cont.copy()
-        sec_cont[:, :, 4:6] = phys.reshape(N_, K_, 2)
-
-    N, K, _ = sec_cont.shape
+    N, K = sec_cont.shape[0], sec_cont.shape[1]
     stick_logits = sec_cont[:, :, 0]  # (N, K)
     dir_local = sec_cont[:, :, 1:4].copy()  # (N, K, 3)
-    log_mass = sec_cont[:, :, 4]  # (N, K)
-    charge = sec_cont[:, :, 5]  # (N, K)
 
     # Flow-matching output isn't guaranteed unit norm; normalise before the
     # rotation below, which preserves magnitude rather than fixing it up.
@@ -666,6 +661,54 @@ def decode_secondaries(
                 pre_dir[valid], dir_local[valid, i]
             )
 
+    return sec_E, sec_dir_world, sec_valid
+
+
+def decode_secondaries(
+    sec_cont: np.ndarray,
+    n_sec: np.ndarray,
+    e_sec: np.ndarray,
+    pre_dir: np.ndarray,
+    sec_phys_normalizer: "Normalizer | None" = None,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Inverse of encode_secondaries: continuous targets → physical secondary attrs.
+
+    `particle_type.target = "physical"` only (the type slice is a raw
+    (log_mass, charge) regression target folded straight into `sec_cont`) —
+    `"onehot"`/`"embedding"` decode through `decode_secondary_cont` +
+    `giant.particles.decode_topn_class`/`decode_embedding_nearest` instead,
+    since their type slice isn't (log_mass, charge) at all. See
+    `decode_secondary_cont`'s docstring for why the two share the energy/
+    direction logic below.
+
+    sec_cont: (N, K_MAX, 6) — [stick_logit, local_dir_x, local_dir_y,
+        local_dir_z, log_mass, charge] (log_mass/charge normalised iff
+        `sec_phys_normalizer` was applied when this was produced — e.g. a
+        raw model prediction; pass the same normalizer here to invert it)
+    n_sec: (N,) integer secondary counts
+    e_sec: (N,) total secondary energy budget [MeV]
+    pre_dir: (N, 3) pre-step world-frame direction
+
+    Returns (sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid) each
+    shape (N, K_MAX). mass/charge are the model's raw predicted physical
+    identity for each secondary, used as-is (no snapping to a discrete PDG
+    code) — see giant/particles.py for the separate, reporting-only
+    nearest-PDG lookup callers may apply on top of this for display/
+    bookkeeping purposes.
+    """
+    if sec_phys_normalizer is not None:
+        N_, K_, _ = sec_cont.shape
+        phys = sec_phys_normalizer.inverse_transform(sec_cont[:, :, 4:6].reshape(-1, 2))
+        sec_cont = sec_cont.copy()
+        sec_cont[:, :, 4:6] = phys.reshape(N_, K_, 2)
+
+    sec_E, sec_dir_world, sec_valid = decode_secondary_cont(
+        sec_cont, n_sec, e_sec, pre_dir
+    )
+
+    log_mass = sec_cont[:, :, 4]  # (N, K)
+    charge = sec_cont[:, :, 5]  # (N, K)
+
     # mass is non-negative by construction (inv_log_transform of a real
     # number is always > 0); clip to 0 for padded/invalid slots rather than
     # leaving a spurious small positive floor from the log inverse.

giant/model/network.py

@@ -982,7 +982,9 @@ class Stage2OneShot(nn.Module):
         super().__init__()
         self.generator_kind = generator
         self.noise_dim = noise_dim
+        self.k_max = k_max
         self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"})
+        self.type_dim = stage2_type_dim(self.particle_type_cfg, particle_cfg["emb_dim"])
         self.cond_enc = ConditionEncoder(
             pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
         )

giant/particles.py

@@ -7,16 +7,28 @@ table involved) for isomer/excited nuclear codes the package's ground-state-only
 nuclide table doesn't cover — confirmed necessary for ~32% of the nuclear codes
 actually present in the multi-material dataset
 (`0932fb02-f2ce-43ca-a4ef-60a2b1221bbc.parquet`).
+
+Also holds the v0.3.0 stage-2 categorical-type rollout decode (§3.3/§8/§11.3
+of docs/v0.3.0-design.md): `decode_topn_class`/`decode_embedding_nearest` turn
+`Stage2Autoregressive`/`Stage2OneShot`'s `"onehot"`/`"embedding"` type
+predictions back into concrete PDG codes, the one place a secondary's
+categorical/continuous type representation is ever discretized (its
+free-running history representation stays unsnapped — see
+`giant/sample.py`'s AR loop).
 """
 
 from __future__ import annotations
 
 from functools import lru_cache
+from typing import TYPE_CHECKING
 
 import numpy as np
 from particle import InvalidParticle, Particle, ParticleNotFound
 from particle import pdgid as _pdgid
 
+if TYPE_CHECKING:
+    from giant.data.loader import TopNMap
+
 # First-pass nuclear mass approximation (A * atomic mass unit); no
 # binding-energy correction. Only used for codes missing from `particle`'s
 # ground-state nuclide table -- ground-state codes get the package's real
@@ -111,3 +123,113 @@ def nearest_known_pdg(mass: np.ndarray, charge: np.ndarray, candidates) -> np.nd
     ) ** 2
     idx = d2.argmin(axis=1)
     return codes[idx]
+
+
+def invert_dense_map(m: dict[int, int]) -> dict[int, int]:
+    """index -> key, inverting a dense, bijective value->index map (`pdg_map`,
+    or a `TopNMap.class_map`'s non-"other" entries — see `decode_topn_class`,
+    which needs a *partial* inverse, not this general one, because its "other"
+    index isn't unique-preimage). `pdg_map` itself is always a true bijection
+    (`giant.data.loader.build_index_maps_from_files` enumerates the vocab), so
+    a plain dict-comprehension inversion is exact here — used for
+    `stage2_model.particle_type.target = "embedding"` decode, whose vocabulary
+    is the full dense `pdg_map`, not a top-N-plus-other map."""
+    return {v: k for k, v in m.items()}
+
+
+def decode_topn_class(
+    class_idx: np.ndarray,
+    topn_map: "TopNMap",
+    n_classes: int,
+    other_policy: str = "sample",
+    rng: np.random.Generator | None = None,
+) -> np.ndarray:
+    """`conditioning.particle.type` / `stage2_model.particle_type.target =
+    "onehot"` inference decode (docs/v0.3.0-design.md §3.3): per-row top-N
+    class index -> concrete PDG code.
+
+    class_idx: int array, any shape, values in `[0, n_classes)`.
+    topn_map: the `TopNMap` (`giant.data.loader.build_pdg_topn_map_from_files`)
+        this class index was built from — `class_map` (PDG -> class, injective
+        except at the shared "other" index) plus `other_members` (the
+        empirical within-"other" distribution, needed for `other_policy =
+        "sample"`/`"modal"`).
+    n_classes: `conditioning.particle.emb_dim` — the class count; the "other"
+        bucket is index `n_classes - 1` by construction
+        (`giant.data.loader._topn_plus_other_map`).
+    other_policy: `"sample"` draws from `other_members`' empirical frequency;
+        `"modal"` always the single most common "other" member; `"drop"`
+        returns PDG `0` for those rows (not a valid PDG code — the caller
+        must treat it as "no secondary", the same convention as
+        `TERM_UNKNOWN_PDG` elsewhere in the rollout driver).
+
+    Every non-"other" class index has a unique inverse (the top `n_classes -
+    1` keys each got their own index in `_topn_plus_other_map`), so those
+    rows decode exactly; only "other" rows need `other_policy`.
+    """
+    other_idx = n_classes - 1
+    inv = np.zeros(n_classes, dtype=np.int64)
+    for pdg, idx in topn_map.class_map.items():
+        if idx != other_idx:
+            inv[idx] = pdg
+
+    flat = np.asarray(class_idx, dtype=np.int64).reshape(-1)
+    out = inv[np.clip(flat, 0, n_classes - 1)]
+
+    other_mask = flat == other_idx
+    n_other = int(other_mask.sum())
+    if n_other:
+        if not topn_map.other_members:
+            raise ValueError(
+                "decode_topn_class: 'other' class predicted but "
+                "topn_map.other_members is empty"
+            )
+        members = np.array(list(topn_map.other_members.keys()), dtype=np.int64)
+        counts = np.array(list(topn_map.other_members.values()), dtype=np.float64)
+        if other_policy == "drop":
+            out[other_mask] = 0
+        elif other_policy == "modal":
+            out[other_mask] = members[counts.argmax()]
+        elif other_policy == "sample":
+            rng = rng if rng is not None else np.random.default_rng()
+            probs = counts / counts.sum()
+            out[other_mask] = rng.choice(members, size=n_other, p=probs)
+        else:
+            raise ValueError(f"unknown other_policy {other_policy!r}")
+
+    return out.reshape(np.asarray(class_idx).shape)
+
+
+def decode_embedding_nearest(
+    vectors: np.ndarray,
+    emb_weight: np.ndarray,
+    idx_to_pdg: dict[int, int],
+) -> tuple[np.ndarray, np.ndarray]:
+    """`stage2_model.particle_type.target = "embedding"` inference decode
+    (docs/v0.3.0-design.md §3.3): L1-nearest row of the conditioning's own
+    particle embedding table, since a generative model's continuous output
+    essentially never lands within float tolerance of a table row (the exact-
+    match form is only valid as a round-trip test assertion, never here).
+
+    vectors: `(..., emb_dim)` raw predicted vectors, any leading shape.
+    emb_weight: `(vocab, emb_dim)` — `ConditionEncoder.pdg_emb.weight`,
+        detached and moved to numpy by the caller. This is the SAME table
+        `particle_type.target = "embedding"` was regressed against
+        (`validate_config` requires `conditioning.particle.type =
+        "embedding"` whenever this target is used — one table, not two).
+    idx_to_pdg: `invert_dense_map(pdg_map)` — embedding row index -> PDG.
+
+    Returns `(pdg, l1_dist)`, both shaped like `vectors.shape[:-1]`. `l1_dist`
+    is the §11.3 diagnostic: a heavy tail means the decoder is emitting
+    vectors off the embedding manifold, the direct analogue of the species-
+    collapse symptom this redesign exists to fix.
+    """
+    emb_dim = vectors.shape[-1]
+    flat = np.asarray(vectors, dtype=np.float64).reshape(-1, emb_dim)
+    table = np.asarray(emb_weight, dtype=np.float64)
+    d = np.abs(flat[:, None, :] - table[None, :, :]).sum(axis=-1)  # (N, vocab)
+    nearest = d.argmin(axis=1)
+    dist = d[np.arange(len(nearest)), nearest]
+    pdg = np.array([idx_to_pdg[int(i)] for i in nearest], dtype=np.int64)
+    lead_shape = vectors.shape[:-1]
+    return pdg.reshape(lead_shape), dist.reshape(lead_shape).astype(np.float32)

giant/rollout.py

@@ -17,7 +17,7 @@ treated as detector leakage and not deposited.
 from __future__ import annotations
 
 from collections import Counter
-from typing import Callable, TypedDict
+from typing import TYPE_CHECKING, Callable, TypedDict
 
 import numpy as np
 import torch
@@ -33,18 +33,165 @@ from giant.data.transforms import (
     Normalizer,
     build_cond_features,
     decode_secondaries,
+    decode_secondary_cont,
     energy_simplex_decode,
     inv_local_frame_rotation,
     inv_log_transform,
     reconstruct_post_pos,
 )
-from giant.particles import nearest_known_pdg, particle_phys_array
-from giant.sample import (
-    sample_flow,
-    sample_secondaries,
-    sample_wgan,
-    sample_secondaries_wgan,
+from giant.particles import (
+    decode_embedding_nearest,
+    decode_topn_class,
+    invert_dense_map,
+    nearest_known_pdg,
+    particle_phys_array,
 )
+from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
+
+if TYPE_CHECKING:
+    from giant.data.loader import TopNMap
+
+
+class L1DistCollector:
+    """Accumulates the §11.3 L1-distance diagnostic across a whole rollout
+    run: the L1 distance between each emitted secondary's raw predicted
+    embedding vector and the nearest table row it snapped to (only
+    meaningful under `particle_type.target = "embedding"` —
+    `giant.particles.decode_embedding_nearest`). A heavy tail means the
+    decoder is emitting vectors off the embedding manifold — the direct
+    analogue of the species-collapse symptom the v0.3.0 redesign exists to
+    fix (docs/v0.3.0-design.md §11.3).
+
+    Not folded into `rollout()`'s own return value (which is shape-typed as
+    step records, see `_RECORD_KEYS`/`RolloutSummary`) — passed in and read
+    back by the caller instead, mirroring the existing `on_chunk` pattern.
+    O(1) memory via a fixed log-spaced histogram rather than raw samples,
+    since a heavy right tail is exactly what this diagnostic watches for.
+    """
+
+    def __init__(self, n_bins: int = 50, lo: float = 1e-3, hi: float = 1e3) -> None:
+        self.n = 0
+        self.total = 0.0
+        self.total_sq = 0.0
+        self.minimum = float("inf")
+        self.maximum = 0.0
+        self.hist_edges = np.geomspace(lo, hi, n_bins + 1)
+        self.hist_counts = np.zeros(n_bins, dtype=np.int64)
+
+    def add(self, dist: np.ndarray, valid: np.ndarray) -> None:
+        vals = np.asarray(dist)[np.asarray(valid)]
+        if vals.size == 0:
+            return
+        self.n += int(vals.size)
+        self.total += float(vals.sum())
+        self.total_sq += float(np.square(vals).sum())
+        self.minimum = min(self.minimum, float(vals.min()))
+        self.maximum = max(self.maximum, float(vals.max()))
+        self.hist_counts += np.histogram(vals, bins=self.hist_edges)[0]
+
+    def summary(self) -> dict | None:
+        """`None` if nothing was ever added (target != "embedding", or a
+        run with zero secondaries) — the caller should omit the diagnostic
+        entirely rather than write a degenerate summary."""
+        if self.n == 0:
+            return None
+        mean = self.total / self.n
+        variance = max(self.total_sq / self.n - mean**2, 0.0)
+        return {
+            "n": self.n,
+            "mean": mean,
+            "std": variance**0.5,
+            "min": self.minimum,
+            "max": self.maximum,
+            "hist_edges": self.hist_edges.tolist(),
+            "hist_counts": self.hist_counts.tolist(),
+        }
+
+
+def decode_secondary_identity(
+    sec_decoder: torch.nn.Module,
+    sec_cont: torch.Tensor,
+    sec_type: torch.Tensor,
+    n_sec_np: np.ndarray,
+    e_sec: np.ndarray,
+    pre_dir: np.ndarray,
+    sec_phys_norm: Normalizer,
+    pdg_map: dict[int, int],
+    pdg_topn_map: "TopNMap | None",
+    other_policy: str,
+    rng: np.random.Generator | None,
+) -> tuple[
+    np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray | None
+]:
+    """Decode Stage 2's raw (sec_cont, sec_type) output into physical
+    secondary attributes, branching on `sec_decoder.particle_type_cfg`
+    (docs/v0.3.0-design.md §3.3):
+
+    - `"physical"`: unchanged v0.2 path — `sec_type` already *is* (log_mass,
+      charge), used as the secondary's identity as-is (no snapping).
+    - `"onehot"`: `sec_type` is per-slot class logits — argmax, then
+      `giant.particles.decode_topn_class` (+ `other_policy`) resolves a
+      concrete PDG, whose real physics (log_mass, charge) then come from
+      `giant.particles.particle_phys_array` — unlike "physical", the PDG
+      resolution IS the secondary's identity here, not just a reporting
+      label.
+    - `"embedding"`: `sec_type` is a raw vector in the conditioning's own
+      embedding space — `giant.particles.decode_embedding_nearest` L1-snaps
+      it to the nearest table row for the PDG (+ physics via
+      `particle_phys_array`), and also returns the L1 distance (§11.3
+      diagnostic — see `giant/rollout.py`'s L1-distance accumulator).
+
+    Returns (sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg,
+    sec_type_l1_dist) — the last is `None` except under `"embedding"`.
+    """
+    target = sec_decoder.particle_type_cfg.get("target", "physical")
+
+    if target == "physical":
+        sec_full = torch.cat([sec_cont, sec_type], dim=-1).cpu().numpy()
+        sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid = decode_secondaries(
+            sec_full, n_sec_np, e_sec, pre_dir, sec_phys_normalizer=sec_phys_norm
+        )
+        sec_pdg = nearest_known_pdg(
+            sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()
+        ).reshape(sec_mass.shape)
+        return sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg, None
+
+    sec_E, sec_dir_world, sec_valid = decode_secondary_cont(
+        sec_cont.cpu().numpy(), n_sec_np, e_sec, pre_dir
+    )
+    sec_type_np = sec_type.cpu().numpy()
+    l1_dist = None
+
+    if target == "onehot":
+        if pdg_topn_map is None:
+            raise RuntimeError(
+                "particle_type.target='onehot' rollout needs pdg_topn_map "
+                "(the checkpoint's saved top-N map) — see ckpt['pdg_topn_map']"
+            )
+        class_idx = sec_type_np.argmax(axis=-1)
+        sec_pdg = decode_topn_class(
+            class_idx,
+            pdg_topn_map,
+            n_classes=sec_decoder.type_dim,
+            other_policy=other_policy,
+            rng=rng,
+        )
+    else:  # "embedding"
+        idx_to_pdg = invert_dense_map(pdg_map)
+        emb_weight = sec_decoder.cond_enc.pdg_emb.weight.detach().cpu().numpy()
+        sec_pdg, l1_dist = decode_embedding_nearest(sec_type_np, emb_weight, idx_to_pdg)
+        l1_dist = np.where(sec_valid, l1_dist, 0.0).astype(np.float32)
+
+    sec_mass, sec_charge = particle_phys_array(sec_pdg.reshape(-1)).T
+    sec_mass = np.where(sec_valid, sec_mass.reshape(sec_pdg.shape), 0.0).astype(
+        np.float32
+    )
+    sec_charge = np.where(sec_valid, sec_charge.reshape(sec_pdg.shape), 0.0).astype(
+        np.float32
+    )
+    sec_pdg = np.where(sec_valid, sec_pdg, 0).astype(np.int64)
+    return sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg, l1_dist
+
 
 # Record columns produced per step / per terminal marker.
 _RECORD_KEYS = [
@@ -311,7 +458,12 @@ def rollout(
     escape_threshold: float | None = None,
     on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None,
     conditioning: str = "embedding",
-    mode: str = "flow",
+    pdg_topn_map: "TopNMap | None" = None,
+    other_policy: str = "sample",
+    seed: int | None = None,
+    stage1_ddpm_steps: int = 1000,
+    stage2_ddpm_steps: int = 1000,
+    l1_dist_collector: "L1DistCollector | None" = None,
 ) -> dict[str, np.ndarray] | RolloutSummary:
     """Run showers to completion.
 
@@ -324,12 +476,30 @@ def rollout(
     dict[str, int]}`. Use this for large `--n-events`/`--max-steps` runs,
     where the full record set would otherwise scale with
     `n_events * max_steps * avg_tracks_per_event`.
+
+    There is no `mode` parameter — each stage's generative objective is read
+    directly off the model instance's own `generator_kind`
+    (docs/v0.3.0-design.md decision 2: stage 1 and stage 2 objectives are
+    independent, e.g. `stage1_model.generator="flow"` +
+    `stage2_model.generator="wgan"`), and the decoder (one-shot vs
+    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
+    `other_policy = "sample"` draw only (torch/numpy sampling itself is
+    seeded by the caller, same as today).
+
+    `l1_dist_collector`, if given, accumulates the §11.3 embedding-distance
+    diagnostic across the whole run — see `L1DistCollector`. Only populated
+    under `particle_type.target = "embedding"`; a no-op otherwise.
     """
     device = device or torch.device("cpu")
     stage1_model.eval()
     sec_decoder.eval()
     if escape_threshold is not None:
         oracle.escape_threshold = float(escape_threshold)
+    rng = np.random.default_rng(seed)
 
     frontier, counts = make_seed_frontier(
         seeds["event_id"],
@@ -365,7 +535,12 @@ def rollout(
                     device,
                     max_tracks_per_event,
                     conditioning,
-                    mode,
+                    pdg_topn_map,
+                    other_policy,
+                    rng,
+                    stage1_ddpm_steps,
+                    stage2_ddpm_steps,
+                    l1_dist_collector,
                 )
             )
         frontier = _concat_frontiers(next_parts)
@@ -396,7 +571,12 @@ def _step_chunk(
     device,
     max_tracks_per_event,
     conditioning,
-    mode="flow",
+    pdg_topn_map,
+    other_policy,
+    rng,
+    stage1_ddpm_steps,
+    stage2_ddpm_steps,
+    l1_dist_collector,
 ) -> dict[str, np.ndarray]:
     """Advance one chunk of tracks by a single step; return the next frontier."""
     n = len(tr["event_id"])
@@ -480,10 +660,9 @@ def _step_chunk(
     cc = torch.from_numpy(cond_cont).float().to(device)
     ck = torch.from_numpy(cond_cat).long().to(device)
 
-    if mode == "wgan":
-        stage1_norm, n_sec_pred = sample_wgan(stage1_model, cc, ck)
-    else:
-        stage1_norm, n_sec_pred = sample_flow(stage1_model, cc, ck, steps=steps)
+    stage1_norm, n_sec_pred_stage1 = sample_stage1(
+        stage1_model, cc, ck, steps, stage1_ddpm_steps
+    )
     raw = tgt_norm.inverse_transform(stage1_norm.cpu().numpy())
 
     step_length = inv_log_transform(raw[:, 0])
@@ -503,32 +682,40 @@ def _step_chunk(
         tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local
     )
 
+    n_sec_pred = resolve_n_sec(
+        stage1_model, sec_decoder, cc, ck, stage1_norm, n_sec_pred_stage1
+    )
     n_sec_np = n_sec_pred.cpu().numpy().astype(np.int64)
 
     # --- Secondaries ---
-    # No snapping: sec_mass/sec_charge are the model's raw predicted physical
-    # identity, used as-is for the spawned track's own future conditioning.
-    # sec_pdg_code below is a *separate*, reporting-only nearest-known-PDG
-    # label (never fed back into the model) — see giant/particles.py.
-    if mode == "wgan":
-        sec_cont, sec_phys, _valid = sample_secondaries_wgan(
-            sec_decoder, cc, ck, stage1_norm, n_sec_pred
-        )
-    else:
-        sec_cont, sec_phys, _valid = sample_secondaries(
-            sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps
-        )
-    sec_full = torch.cat([sec_cont, sec_phys], dim=-1).cpu().numpy()
-    sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid = decode_secondaries(
-        sec_full,
-        n_sec_np,
-        e_sec,
-        tr["pre_dir"],
-        sec_phys_normalizer=sec_phys_norm,
+    # No snapping for "physical"/history-facing state elsewhere in the
+    # pipeline: sec_mass/sec_charge (or, for "onehot"/"embedding", the
+    # resolved sec_pdg -> real physics) are the secondary's identity, used
+    # as-is for the spawned track's own future conditioning — see
+    # decode_secondary_identity's docstring for how each
+    # particle_type.target differs on whether PDG resolution is a real
+    # identity decision or just a reporting label.
+    sec_cont, sec_type, _valid = sample_stage2(
+        sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps
     )
-    sec_pdg_code = nearest_known_pdg(
-        sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()
-    ).reshape(sec_mass.shape)
+    sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, sec_type_l1_dist = (
+        decode_secondary_identity(
+            sec_decoder,
+            sec_cont,
+            sec_type,
+            n_sec_np,
+            e_sec,
+            tr["pre_dir"],
+            sec_phys_norm,
+            pdg_map,
+            pdg_topn_map,
+            other_policy,
+            rng,
+        )
+    )
+    sec_valid = np.arange(sec_E.shape[1])[None, :] < n_sec_np[:, None]
+    if l1_dist_collector is not None and sec_type_l1_dist is not None:
+        l1_dist_collector.add(sec_type_l1_dist, sec_valid)
 
     edep = edep.astype(np.float64)
     post_E = post_E.astype(np.float64)

giant/sample.py

@@ -1,26 +1,23 @@
 import torch
+import torch.nn.functional as F
 
-from giant.constants import K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM
+from giant.constants import CONT_SLOT_DIM, X_DIM
+from giant.model.network import Stage2Autoregressive, stage2_trunk_sec_dim
+from giant.model.schedule import CosineSchedule
 
 
-def _slots_from_flat(
-    x: torch.Tensor, n_sec_pred: torch.Tensor
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Reshape a flat (B, SEC_DIM) decoder output into per-slot tensors.
-
-    Returns (sec_cont, sec_phys, sec_valid) — see `sample_secondaries`'s
-    docstring for their shapes/meaning. Shared by both the flow-matching and
-    WGAN Stage-2 samplers, which differ only in how `x` was produced.
-    """
-    B = x.size(0)
-    device = x.device
-    x_slots = x.view(B, K_MAX, SEC_SLOT_DIM)
-    sec_cont = x_slots[:, :, :4]
-    sec_phys = x_slots[:, :, 4:]
-    sec_valid = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(
-        1
-    )
-    return sec_cont, sec_phys, sec_valid
+def _predict_n_sec_if_owned(
+    model: torch.nn.Module, cond_cont: torch.Tensor, cond_cat: torch.Tensor
+) -> torch.Tensor | None:
+    """Stage-1 `n_sec_head` is only present on a migrated v0.2 checkpoint
+    (docs/v0.3.0-design.md decision 1 moves it to stage 2 for fresh runs —
+    see `Stage1Model`'s docstring). `None` here means "ask stage 2 instead",
+    which every caller (`giant/rollout.py`, `giant/cli.py`) must do for a
+    fresh checkpoint."""
+    if getattr(model, "n_sec_head", None) is None:
+        return None
+    logits = model.predict_n_sec(cond_cont, cond_cat)
+    return logits.argmax(dim=-1)
 
 
 @torch.no_grad()
@@ -29,12 +26,14 @@ def sample_flow(
     cond_cont: torch.Tensor,
     cond_cat: torch.Tensor,
     steps: int = 10,
-) -> tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor | None]:
     """Euler integration of the Stage-1 vector field from t=0 to t=1.
 
     Returns (primary_sample, n_sec_pred):
         primary_sample: (B, X_DIM) — normalised 9D primary post-step output
-        n_sec_pred: (B,) int64 — predicted secondary count
+        n_sec_pred: (B,) int64 — predicted secondary count, or `None` if
+            `model` has no `n_sec_head` (a fresh v0.3.0 Stage1Model — see
+            `_predict_n_sec_if_owned`).
     """
     model.eval()
     B = cond_cont.size(0)
@@ -43,11 +42,142 @@ def sample_flow(
     dt = 1.0 / steps
     for i in range(steps):
         t = torch.full((B,), i * dt, device=device)
-        v = model(x, t, cond_cont, cond_cat)
+        v = model(x, cond_cont, cond_cat, t=t)
         x = x + v * dt
-    n_sec_logits = model.predict_n_sec(cond_cont, cond_cat)
-    n_sec_pred = n_sec_logits.argmax(dim=-1)
-    return x, n_sec_pred
+    return x, _predict_n_sec_if_owned(model, cond_cont, cond_cat)
+
+
+@torch.no_grad()
+def sample_ddpm(
+    model: torch.nn.Module,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+    schedule,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Full DDPM ancestral sampling (T reverse steps). Returns (sample, n_sec_pred)
+    — see `sample_flow`'s docstring for the `n_sec_pred` `None` case."""
+    model.eval()
+    B = cond_cont.size(0)
+    device = cond_cont.device
+    x = torch.randn(B, X_DIM, device=device)
+    T = schedule.T
+    for i in reversed(range(T)):
+        t_norm = torch.full((B,), i / T, device=device)
+        eps_pred = model(x, cond_cont, cond_cat, t=t_norm)
+        beta = schedule.betas[i]
+        alpha = schedule.alphas[i]
+        alpha_bar = schedule.alpha_bars[i]
+        z = torch.randn_like(x) if i > 0 else torch.zeros_like(x)
+        x = (1.0 / alpha.sqrt()) * (
+            x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred
+        ) + beta.sqrt() * z
+    return x, _predict_n_sec_if_owned(model, cond_cont, cond_cat)
+
+
+@torch.no_grad()
+def sample_ddim(
+    model: torch.nn.Module,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+    schedule,
+    steps: int = 50,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """DDIM deterministic sampling (Song et al. 2020). Returns (sample, n_sec_pred)
+    — see `sample_flow`'s docstring for the `n_sec_pred` `None` case."""
+    model.eval()
+    B = cond_cont.size(0)
+    device = cond_cont.device
+    T = schedule.T
+    timesteps = torch.linspace(T - 1, 0, steps, dtype=torch.long, device=device)
+    x = torch.randn(B, X_DIM, device=device)
+    for step_idx, ts in enumerate(timesteps):
+        t_idx = int(ts.item())
+        t_norm = torch.full((B,), t_idx / T, device=device)
+        eps_pred = model(x, cond_cont, cond_cat, t=t_norm)
+        ab_t = schedule.alpha_bars[t_idx]
+        if step_idx + 1 < len(timesteps):
+            ab_prev = schedule.alpha_bars[int(timesteps[step_idx + 1].item())]
+        else:
+            ab_prev = torch.ones(1, device=device)
+        x0_pred = (x - (1.0 - ab_t).sqrt() * eps_pred) / ab_t.sqrt()
+        x = ab_prev.sqrt() * x0_pred + (1.0 - ab_prev).sqrt() * eps_pred
+    return x, _predict_n_sec_if_owned(model, cond_cont, cond_cat)
+
+
+@torch.no_grad()
+def sample_wgan(
+    generator: torch.nn.Module,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Single-pass Stage-1 WGAN generator sample. Returns (sample, n_sec_pred)
+    — see `sample_flow`'s docstring for the `n_sec_pred` `None` case."""
+    generator.eval()
+    B = cond_cont.size(0)
+    z = torch.randn(B, generator.noise_dim, device=cond_cont.device)
+    x = generator(z, cond_cont, cond_cat)
+    return x, _predict_n_sec_if_owned(generator, cond_cont, cond_cat)
+
+
+def _stage2_flat_width(sec_decoder: torch.nn.Module) -> int:
+    """The width of `sec_decoder`'s own trunk in/out vector — folded
+    (continuous + type) under `particle_type.target = "physical"` or
+    `generator = "wgan"`, continuous-only otherwise (the type slice then
+    comes from `predict_type` instead — see `stage2_trunk_sec_dim`'s
+    docstring, docs/v0.3.0-design.md decision 2)."""
+    return stage2_trunk_sec_dim(
+        sec_decoder.particle_type_cfg,
+        sec_decoder.generator_kind,
+        sec_decoder.k_max,
+        sec_decoder.type_dim,
+    )
+
+
+def _type_folded(sec_decoder: torch.nn.Module) -> bool:
+    target = sec_decoder.particle_type_cfg.get("target", "physical")
+    return target == "physical" or sec_decoder.generator_kind == "wgan"
+
+
+def _decode_stage2_flat(
+    sec_decoder: torch.nn.Module,
+    x: torch.Tensor,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+    stage1_out: torch.Tensor,
+    n_sec_pred: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Reshape a flat `(B, flat_width)` `Stage2OneShot` output into per-slot
+    tensors, generator/`particle_type.target`-agnostic (docs/v0.3.0-design.md
+    decision 2/3): shared by `sample_secondaries`/`sample_secondaries_wgan`,
+    which differ only in how `x` was produced.
+
+    Returns (sec_cont, sec_type, sec_valid):
+        sec_cont:  (B, k_max, CONT_SLOT_DIM) — [stick_logit, local_dir]
+        sec_type:  (B, k_max, type_dim) — under `target="physical"` this is
+            [log_mass, charge] (normalised iff the checkpoint's sec_phys
+            normalizer was applied at training time — denormalize before
+            treating as physical units; see
+            giant.data.transforms.decode_secondaries); under `"onehot"` /
+            `"embedding"` it is raw class logits / an embedding-space vector
+            — decode via giant.particles.decode_topn_class /
+            decode_embedding_nearest (see giant/rollout.py).
+        sec_valid: (B, k_max) bool — True for slots i < n_sec_pred
+    """
+    B = x.size(0)
+    device = x.device
+    k_max = sec_decoder.k_max
+    type_dim = sec_decoder.type_dim
+    if _type_folded(sec_decoder):
+        x_slots = x.view(B, k_max, CONT_SLOT_DIM + type_dim)
+        sec_cont = x_slots[:, :, :CONT_SLOT_DIM]
+        sec_type = x_slots[:, :, CONT_SLOT_DIM:]
+    else:
+        sec_cont = x.view(B, k_max, CONT_SLOT_DIM)
+        sec_type = sec_decoder.predict_type(cond_cont, cond_cat, stage1_out)
+    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(
+        1
+    )
+    return sec_cont, sec_type, sec_valid
 
 
 @torch.no_grad()
@@ -59,76 +189,27 @@ def sample_secondaries(
     n_sec_pred: torch.Tensor,
     steps: int = 10,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Euler integration of the Stage-2 vector field; return raw slot outputs.
+    """Euler integration of `Stage2OneShot`'s (flow/ddpm) vector field; return
+    raw slot outputs — see `_decode_stage2_flat`'s docstring for the returned
+    (sec_cont, sec_type, sec_valid) shapes/meaning.
 
     n_sec_pred: (B,) int64 — number of valid secondaries per step
-
-    Returns (sec_cont, sec_phys, sec_valid):
-        sec_cont:  (B, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z]
-        sec_phys:  (B, K_MAX, PARTICLE_PHYS_DIM) — predicted [log_mass, charge]
-                   per slot (normalised iff the checkpoint's sec_phys
-                   normalizer was applied at training time — denormalize
-                   before treating as physical units; see
-                   giant.data.transforms.decode_secondaries). Used as-is —
-                   no snapping to a discrete PDG code.
-        sec_valid: (B, K_MAX) bool — True for slots i < n_sec_pred
     """
     sec_decoder.eval()
     B = cond_cont.size(0)
     device = cond_cont.device
+    flat_width = _stage2_flat_width(sec_decoder)
 
-    x = torch.randn(B, SEC_DIM, device=device)
+    x = torch.randn(B, flat_width, device=device)
     dt = 1.0 / steps
     for i in range(steps):
         t = torch.full((B,), i * dt, device=device)
-        v = sec_decoder(x, t, cond_cont, cond_cat, stage1_out)
+        v = sec_decoder(x, cond_cont, cond_cat, stage1_out, t=t)
         x = x + v * dt
 
-    return _slots_from_flat(x, n_sec_pred)
-
-
-@torch.no_grad()
-def sample_ddpm(
-    model: torch.nn.Module,
-    cond_cont: torch.Tensor,
-    cond_cat: torch.Tensor,
-    schedule,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Full DDPM ancestral sampling (T reverse steps). Returns (sample, n_sec_pred)."""
-    model.eval()
-    B = cond_cont.size(0)
-    device = cond_cont.device
-    x = torch.randn(B, X_DIM, device=device)
-    T = schedule.T
-    for i in reversed(range(T)):
-        t_norm = torch.full((B,), i / T, device=device)
-        eps_pred = model(x, t_norm, cond_cont, cond_cat)
-        beta = schedule.betas[i]
-        alpha = schedule.alphas[i]
-        alpha_bar = schedule.alpha_bars[i]
-        z = torch.randn_like(x) if i > 0 else torch.zeros_like(x)
-        x = (1.0 / alpha.sqrt()) * (
-            x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred
-        ) + beta.sqrt() * z
-    n_sec_logits = model.predict_n_sec(cond_cont, cond_cat)
-    n_sec_pred = n_sec_logits.argmax(dim=-1)
-    return x, n_sec_pred
-
-
-@torch.no_grad()
-def sample_wgan(
-    generator: torch.nn.Module,
-    cond_cont: torch.Tensor,
-    cond_cat: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Single-pass Stage-1 WGAN generator sample. Returns (sample, n_sec_pred)."""
-    generator.eval()
-    B = cond_cont.size(0)
-    z = torch.randn(B, generator.noise_dim, device=cond_cont.device)
-    x = generator(z, cond_cont, cond_cat)
-    n_sec_logits = generator.predict_n_sec(cond_cont, cond_cat)
-    n_sec_pred = n_sec_logits.argmax(dim=-1)
-    return x, n_sec_pred
+    return _decode_stage2_flat(
+        sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred
+    )
 
 
 @torch.no_grad()
@@ -139,41 +220,231 @@ def sample_secondaries_wgan(
     stage1_out: torch.Tensor,
     n_sec_pred: torch.Tensor,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Single-pass Stage-2 WGAN generator sample; see `sample_secondaries`'s
-    docstring for the returned (sec_cont, sec_phys, sec_valid) shapes."""
+    """Single-pass `Stage2OneShot` WGAN generator sample; see
+    `_decode_stage2_flat`'s docstring for the returned (sec_cont, sec_type,
+    sec_valid) shapes/meaning."""
     sec_decoder.eval()
     B = cond_cont.size(0)
     z = torch.randn(B, sec_decoder.noise_dim, device=cond_cont.device)
     x = sec_decoder(z, cond_cont, cond_cat, stage1_out)
-    return _slots_from_flat(x, n_sec_pred)
+    return _decode_stage2_flat(
+        sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred
+    )
 
 
 @torch.no_grad()
-def sample_ddim(
-    model: torch.nn.Module,
+def sample_secondaries_ar(
+    sec_decoder: torch.nn.Module,
     cond_cont: torch.Tensor,
     cond_cat: torch.Tensor,
-    schedule,
-    steps: int = 50,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """DDIM deterministic sampling (Song et al. 2020). Returns (sample, n_sec_pred)."""
-    model.eval()
+    stage1_out: torch.Tensor,
+    n_sec_pred: torch.Tensor,
+    steps: int = 10,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """`Stage2Autoregressive` inference loop (docs/v0.3.0-design.md §6.4):
+    one token at a time, in descending-energy slot order, `k_max` sequential
+    calls. Unlike training (teacher forcing, §6.2 point 3 — a single
+    parallel pass over ground-truth tokens, see
+    `giant.train._assemble_stage2_ar_inputs`), there is no ground truth at
+    inference: each token's conditioning is built free-running, from the
+    PREVIOUS TOKEN'S OWN just-generated output — the train/inference gap
+    §6.2 point 4 explicitly flags as the cost of markov history's
+    expressiveness.
+
+    A `{flow,ddpm}` token costs `steps` ODE substeps; `wgan` costs one pass —
+    §6.4's "K sequential forwards" cost note applies per-token here, not
+    once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per
+    physics step.
+
+    The free-running history feature stays UNSNAPPED (mirrors the
+    established "no snapping" precedent for `particle_type.target =
+    "physical"` secondaries feeding their own future conditioning):
+    `"physical"` carries the raw (log_mass, charge) forward as-is;
+    `"embedding"` carries the raw predicted vector as-is; `"onehot"` is the
+    one exception — its history slot must be a probability-simplex-shaped
+    vector (that's what `MarkovHistory` was trained on, `_type_repr`'s
+    `F.one_hot` ground truth), so it's the hard one-hot of `argmax(logits)`,
+    not the raw logits themselves. Discretizing further, into a concrete PDG
+    code, only ever happens once — at secondary-spawn time in
+    `giant/rollout.py` — never inside this loop.
+
+    Returns (sec_cont, sec_type, sec_valid) — same shapes/meaning as
+    `sample_secondaries`/`sample_secondaries_wgan`'s (see
+    `_decode_stage2_flat`'s docstring); `sec_type` is raw per-slot output in
+    all three `particle_type.target` cases (never one-hot-collapsed), so the
+    caller decodes it exactly the same way regardless of which decoder
+    produced it.
+    """
+    sec_decoder.eval()
     B = cond_cont.size(0)
     device = cond_cont.device
-    T = schedule.T
-    timesteps = torch.linspace(T - 1, 0, steps, dtype=torch.long, device=device)
-    x = torch.randn(B, X_DIM, device=device)
-    for step_idx, ts in enumerate(timesteps):
-        t_idx = int(ts.item())
-        t_norm = torch.full((B,), t_idx / T, device=device)
-        eps_pred = model(x, t_norm, cond_cont, cond_cat)
-        ab_t = schedule.alpha_bars[t_idx]
-        if step_idx + 1 < len(timesteps):
-            ab_prev = schedule.alpha_bars[int(timesteps[step_idx + 1].item())]
+    k_max = sec_decoder.k_max
+    type_dim = sec_decoder.type_dim
+    generator = sec_decoder.generator_kind
+    target = sec_decoder.particle_type_cfg.get("target", "physical")
+    type_folded = _type_folded(sec_decoder)
+    token_dim = CONT_SLOT_DIM + type_dim if type_folded else CONT_SLOT_DIM
+
+    sec_cont = torch.zeros(B, k_max, CONT_SLOT_DIM, device=device)
+    sec_type = torch.zeros(B, k_max, type_dim, device=device)
+
+    # Running per-token state, threaded from one slot to the next.
+    prev_repr = torch.zeros(B, CONT_SLOT_DIM + type_dim, device=device)
+    remaining = torch.ones(B, device=device)
+
+    for k in range(k_max):
+        has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device)
+        history_feat = prev_repr.unsqueeze(1)  # (B, 1, CONT_SLOT_DIM + type_dim)
+        remaining_frac = remaining.unsqueeze(1)  # (B, 1)
+        slot_idx = torch.full(
+            (B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32
+        )
+
+        if generator == "wgan":
+            z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
+            token = sec_decoder(
+                z,
+                cond_cont,
+                cond_cat,
+                stage1_out,
+                history_feat,
+                has_prev,
+                remaining_frac,
+                slot_idx,
+            )
         else:
-            ab_prev = torch.ones(1, device=device)
-        x0_pred = (x - (1.0 - ab_t).sqrt() * eps_pred) / ab_t.sqrt()
-        x = ab_prev.sqrt() * x0_pred + (1.0 - ab_prev).sqrt() * eps_pred
-    n_sec_logits = model.predict_n_sec(cond_cont, cond_cat)
-    n_sec_pred = n_sec_logits.argmax(dim=-1)
-    return x, n_sec_pred
+            x = torch.randn(B, 1, token_dim, device=device)
+            dt = 1.0 / steps
+            for i in range(steps):
+                t = torch.full((B, 1), i * dt, device=device)
+                v = sec_decoder(
+                    x,
+                    cond_cont,
+                    cond_cat,
+                    stage1_out,
+                    history_feat,
+                    has_prev,
+                    remaining_frac,
+                    slot_idx,
+                    t=t,
+                )
+                x = x + v * dt
+            token = x
+
+        token = token.squeeze(1)  # (B, token_dim)
+        cont_k = token[:, :CONT_SLOT_DIM]
+        if type_folded:
+            type_k = token[:, CONT_SLOT_DIM:]
+        else:
+            type_k = sec_decoder.predict_type(
+                cond_cont,
+                cond_cat,
+                stage1_out,
+                history_feat,
+                has_prev,
+                remaining_frac,
+                slot_idx,
+            ).squeeze(1)
+
+        sec_cont[:, k] = cont_k
+        sec_type[:, k] = type_k
+
+        if target == "onehot":
+            type_for_history = F.one_hot(
+                type_k.argmax(dim=-1), num_classes=type_dim
+            ).float()
+        else:
+            type_for_history = type_k
+
+        stick_fraction = torch.sigmoid(cont_k[:, 0])
+        prev_repr = torch.cat(
+            [stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1
+        )
+        remaining = torch.clamp(remaining * (1.0 - stick_fraction), min=0.0)
+
+    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(
+        1
+    )
+    return sec_cont, sec_type, sec_valid
+
+
+# ---------------------------------------------------------------------------
+# Per-stage dispatch — shared by giant/rollout.py and giant/cli.py's
+# `predict` command, since both need "given a stage model, produce a
+# sample" without hand-picking the sampler themselves (docs/v0.3.0-design.md
+# decision 2: each stage's generative objective is independent, read off the
+# model's own `generator_kind`, not a caller-supplied `mode` string).
+# ---------------------------------------------------------------------------
+
+
+def sample_stage1(
+    stage1_model: torch.nn.Module,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+    steps: int,
+    ddpm_steps: int = 1000,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Dispatches on `stage1_model.generator_kind`."""
+    kind = stage1_model.generator_kind
+    if kind == "wgan":
+        return sample_wgan(stage1_model, cond_cont, cond_cat)
+    if kind == "ddpm":
+        schedule = CosineSchedule(T=ddpm_steps).to(cond_cont.device)
+        return sample_ddpm(stage1_model, cond_cont, cond_cat, schedule)
+    return sample_flow(stage1_model, cond_cont, cond_cat, steps=steps)
+
+
+def sample_stage2(
+    sec_decoder: torch.nn.Module,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+    stage1_out: torch.Tensor,
+    n_sec_pred: torch.Tensor,
+    steps: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Dispatches on `decoder` (one-shot vs autoregressive — the class
+    itself, via `isinstance`) and `sec_decoder.generator_kind` (flow/ddpm/
+    wgan). DDPM secondaries aren't supported — no `Stage2*` class was ever
+    built with `generator="ddpm"` in practice and `flow_matching_loss_secondary*`
+    is the only stage-2 training path that exists for the non-adversarial
+    case, so there's nothing to dispatch to here.
+    """
+    if isinstance(sec_decoder, Stage2Autoregressive):
+        return sample_secondaries_ar(
+            sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps
+        )
+    if sec_decoder.generator_kind == "wgan":
+        return sample_secondaries_wgan(
+            sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred
+        )
+    return sample_secondaries(
+        sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps
+    )
+
+
+def resolve_n_sec(
+    stage1_model: torch.nn.Module,
+    sec_decoder: torch.nn.Module,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+    stage1_out: torch.Tensor,
+    n_sec_pred: torch.Tensor | None,
+) -> torch.Tensor:
+    """`n_sec_pred` is already populated when `stage1_model` owns a legacy
+    `n_sec_head` (a migrated v0.2 checkpoint — see `Stage1Model`'s
+    docstring); otherwise ask stage 2, which owns it by default under
+    decision 1 (docs/v0.3.0-design.md §2). Raises if neither stage owns a
+    head at all — the only way that happens is `stage2_model.n_sec.mode`
+    other than `"head"` (`"truth"`/`"stop_token"`), neither of which is a
+    valid rollout-/predict-capable checkpoint (§3.3, §9)."""
+    if n_sec_pred is not None:
+        return n_sec_pred
+    if getattr(sec_decoder, "n_sec_head", None) is None:
+        raise RuntimeError(
+            "checkpoint has no n_sec_head on either stage — needs "
+            "stage2_model.n_sec.mode = 'head' (the default); 'truth' is "
+            "standalone-evaluation-only and 'stop_token' isn't implemented "
+            "(docs/v0.3.0-design.md §3.3/§9)"
+        )
+    logits = sec_decoder.predict_n_sec(cond_cont, cond_cat, stage1_out)
+    return logits.argmax(dim=-1)

giant/train.py

@@ -991,27 +991,31 @@ _WARNED_MARGINAL_VALIDATION_BROKEN = False
 
 
 def _try_validate_marginals(trainer: StageTrainer, val_loader, device, **kwargs):
-    """`validate_marginals` delegates to `giant.sample`'s samplers, which
-    still assume Stage 1 always owns `n_sec_head` (v0.2 behaviour) — under
-    decision 1 (docs/v0.3.0-design.md §2) a fresh Stage2OneShot owns it by
-    default instead, so this currently raises for any non-legacy checkpoint.
-    `giant/sample.py` needs a per-stage generator/n_sec update (design doc
-    §10, deferred to step 6); until then this degrades gracefully with a
-    one-time warning instead of crashing the whole training run.
+    """`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
     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/sample.py's per-stage n_sec update lands.
+        # 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/sample.py doesn't yet "
-                "support a Stage2-owned n_sec head (docs/v0.3.0-design.md "
-                "step 6); val_marginal_kl stays NaN, and wgan best-checkpoint "
-                "selection falls back to the Wasserstein-distance magnitude.",
+                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

tests/test_cli_predict.py

@@ -6,14 +6,14 @@ import yaml
 
 from giant.cli import (
     _CEPH_PREDICTIONS,
-    _check_v030_onehot_support,
+    _check_conditioning_onehot_support,
     _resolve_prediction_output,
     _write_prediction_ref,
 )
 
 
 # ---------------------------------------------------------------------------
-# _check_v030_onehot_support
+# _check_conditioning_onehot_support
 # ---------------------------------------------------------------------------
 
 
@@ -29,42 +29,43 @@ def _nested_model_cfg(
     }
 
 
-def test_check_v030_onehot_support_allows_physical():
-    _check_v030_onehot_support(_nested_model_cfg(), "predict")  # no raise
+def test_check_conditioning_onehot_support_allows_physical():
+    _check_conditioning_onehot_support(_nested_model_cfg(), "predict")  # no raise
 
 
-def test_check_v030_onehot_support_rejects_onehot_particle_conditioning():
+def test_check_conditioning_onehot_support_rejects_onehot_particle_conditioning():
     cfg = _nested_model_cfg(particle_type="onehot")
     with pytest.raises(typer.Exit):
-        _check_v030_onehot_support(cfg, "predict")
+        _check_conditioning_onehot_support(cfg, "predict")
 
 
-def test_check_v030_onehot_support_rejects_onehot_material_conditioning():
+def test_check_conditioning_onehot_support_rejects_onehot_material_conditioning():
     cfg = _nested_model_cfg(material_type="onehot")
     with pytest.raises(typer.Exit):
-        _check_v030_onehot_support(cfg, "rollout")
+        _check_conditioning_onehot_support(cfg, "rollout")
 
 
-def test_check_v030_onehot_support_rejects_onehot_particle_type_target():
+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")
-    with pytest.raises(typer.Exit):
-        _check_v030_onehot_support(cfg, "predict")
+    _check_conditioning_onehot_support(cfg, "predict")  # no raise
 
 
-def test_check_v030_onehot_support_rejects_embedding_particle_type_target():
+def test_check_conditioning_onehot_support_allows_embedding_particle_type_target():
     cfg = _nested_model_cfg(
         particle_type="embedding", material_type="embedding", target="embedding"
     )
-    with pytest.raises(typer.Exit):
-        _check_v030_onehot_support(cfg, "predict")
+    _check_conditioning_onehot_support(cfg, "predict")  # no raise
 
 
-def test_check_v030_onehot_support_is_noop_for_v02_flat_model_config():
+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/embedding-target, so this must be a
-    silent no-op rather than crash on `.get("particle")` against a string."""
+    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_v030_onehot_support(cfg, "predict")  # no raise
+    _check_conditioning_onehot_support(cfg, "predict")  # no raise
 
 
 # ---------------------------------------------------------------------------

tests/test_flow.py

@@ -1,4 +1,3 @@
-import pytest
 import torch
 from giant.constants import COND_DIM
 from giant.model.network import Stage1Model
@@ -8,13 +7,6 @@ from giant.sample import sample_flow, sample_ddim
 PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
 MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
 
-_SAMPLE_XFAIL_REASON = (
-    "giant/sample.py isn't updated yet — its sample_flow/sample_ddim call "
-    "models positionally as model(x, t, cond_cont, cond_cat), which doesn't "
-    "match Stage1Model's new forward signature. Deferred to "
-    "docs/v0.3.0-design.md step 6."
-)
-
 
 def _small_model():
     return Stage1Model(
@@ -54,7 +46,6 @@ def test_flow_matching_loss_has_grad():
     assert any(p.grad is not None for p in model.parameters())
 
 
-@pytest.mark.xfail(reason=_SAMPLE_XFAIL_REASON, strict=False)
 def test_sample_flow_shape():
     B = 6
     cond_cont = torch.randn(B, COND_DIM)
@@ -71,7 +62,6 @@ def test_ddpm_loss_nonneg():
     assert loss.item() >= 0.0
 
 
-@pytest.mark.xfail(reason=_SAMPLE_XFAIL_REASON, strict=False)
 def test_sample_ddim_shape():
     B = 4
     schedule = CosineSchedule(T=50)

tests/test_particles.py

@@ -1,7 +1,11 @@
 import numpy as np
 import pytest
 
+from giant.data.loader import TopNMap
 from giant.particles import (
+    decode_embedding_nearest,
+    decode_topn_class,
+    invert_dense_map,
     nearest_known_pdg,
     particle_mass_charge,
     particle_phys_array,
@@ -126,3 +130,101 @@ def test_nearest_known_pdg_shape():
     )
     assert result.shape == (n,)
     assert set(result.tolist()) <= set(candidates)
+
+
+# ── invert_dense_map ─────────────────────────────────────────────────────
+
+
+def test_invert_dense_map_round_trips():
+    pdg_map = {22: 0, 11: 1, -11: 2, 2212: 3}
+    inv = invert_dense_map(pdg_map)
+    for pdg, idx in pdg_map.items():
+        assert inv[idx] == pdg
+
+
+# ── decode_topn_class ────────────────────────────────────────────────────
+
+
+def _topn_fixture():
+    # n_classes=4: photon/electron/positron get their own class (0,1,2),
+    # everything else (proton, neutron) falls into "other" (class 3).
+    class_map = {22: 0, 11: 1, -11: 2, 2212: 3, 2112: 3}
+    other_members = {2212: 7, 2112: 3}
+    return TopNMap(class_map=class_map, other_members=other_members), 4
+
+
+def test_decode_topn_class_known_classes_are_exact():
+    topn_map, n_classes = _topn_fixture()
+    out = decode_topn_class(np.array([0, 1, 2]), topn_map, n_classes)
+    np.testing.assert_array_equal(out, [22, 11, -11])
+
+
+def test_decode_topn_class_other_modal_picks_most_frequent():
+    topn_map, n_classes = _topn_fixture()
+    out = decode_topn_class(np.array([3, 3]), topn_map, n_classes, other_policy="modal")
+    assert (out == 2212).all()  # count 7 > 3
+
+
+def test_decode_topn_class_other_drop_returns_zero_sentinel():
+    topn_map, n_classes = _topn_fixture()
+    out = decode_topn_class(np.array([3]), topn_map, n_classes, other_policy="drop")
+    assert out[0] == 0
+
+
+def test_decode_topn_class_other_sample_stays_within_members():
+    topn_map, n_classes = _topn_fixture()
+    rng = np.random.default_rng(0)
+    out = decode_topn_class(
+        np.full(50, 3), topn_map, n_classes, other_policy="sample", rng=rng
+    )
+    assert set(out.tolist()) <= {2212, 2112}
+
+
+def test_decode_topn_class_unknown_other_policy_raises():
+    topn_map, n_classes = _topn_fixture()
+    with pytest.raises(ValueError):
+        decode_topn_class(np.array([3]), topn_map, n_classes, other_policy="bogus")
+
+
+def test_decode_topn_class_empty_other_members_raises():
+    class_map = {22: 0, 11: 1}
+    topn_map = TopNMap(class_map=class_map, other_members={})
+    with pytest.raises(ValueError):
+        decode_topn_class(np.array([1]), topn_map, 2, other_policy="sample")
+
+
+def test_decode_topn_class_preserves_shape():
+    topn_map, n_classes = _topn_fixture()
+    idx = np.array([[0, 1], [2, 3]])
+    out = decode_topn_class(idx, topn_map, n_classes, other_policy="modal")
+    assert out.shape == (2, 2)
+
+
+# ── decode_embedding_nearest ─────────────────────────────────────────────
+
+
+def test_decode_embedding_nearest_exact_row_recovers_pdg():
+    emb_weight = np.array([[1.0, 0.0], [0.0, 1.0], [-1.0, -1.0]])
+    idx_to_pdg = {0: 22, 1: 11, 2: 2212}
+    vectors = np.array([[0.0, 1.0], [-1.0, -1.0]])  # exact rows 1, 2
+    pdg, dist = decode_embedding_nearest(vectors, emb_weight, idx_to_pdg)
+    np.testing.assert_array_equal(pdg, [11, 2212])
+    np.testing.assert_allclose(dist, [0.0, 0.0], atol=1e-8)
+
+
+def test_decode_embedding_nearest_off_manifold_snaps_to_closest_row():
+    emb_weight = np.array([[1.0, 0.0], [0.0, 1.0]])
+    idx_to_pdg = {0: 22, 1: 11}
+    vectors = np.array([[0.9, 0.2]])  # closer to row 0
+    pdg, dist = decode_embedding_nearest(vectors, emb_weight, idx_to_pdg)
+    assert pdg[0] == 22
+    assert dist[0] > 0.0
+
+
+def test_decode_embedding_nearest_preserves_leading_shape():
+    emb_weight = np.array([[1.0, 0.0], [0.0, 1.0]])
+    idx_to_pdg = {0: 22, 1: 11}
+    vectors = np.random.default_rng(0).standard_normal((3, 4, 2))
+    pdg, dist = decode_embedding_nearest(vectors, emb_weight, idx_to_pdg)
+    assert pdg.shape == (3, 4)
+    assert dist.shape == (3, 4)

tests/test_phase2.py

@@ -19,13 +19,6 @@ from giant.model.schedule import (
 )
 from giant.sample import sample_secondaries
 
-_SAMPLE_SECONDARIES_XFAIL_REASON = (
-    "giant/sample.py isn't updated yet — sample_secondaries calls the "
-    "decoder positionally as decoder(x, t, cond_cont, cond_cat, stage1_out), "
-    "which doesn't match Stage2OneShot's new forward signature. Deferred to "
-    "docs/v0.3.0-design.md step 6."
-)
-
 
 # ── helpers ──────────────────────────────────────────────────────────────────
 
@@ -300,7 +293,6 @@ def test_flow_matching_loss_secondary_ar_has_grad():
 # ── sampling ──────────────────────────────────────────────────────────────────
 
 
-@pytest.mark.xfail(reason=_SAMPLE_SECONDARIES_XFAIL_REASON, strict=False)
 def test_sample_secondaries_shapes():
     B, pdg, mat = 6, 3, 2
     decoder = _sec_decoder(pdg, mat)
@@ -316,7 +308,6 @@ def test_sample_secondaries_shapes():
     assert sec_valid.dtype == torch.bool
 
 
-@pytest.mark.xfail(reason=_SAMPLE_SECONDARIES_XFAIL_REASON, strict=False)
 def test_sample_secondaries_valid_mask_matches_n_sec():
     B, pdg, mat = 4, 3, 2
     decoder = _sec_decoder(pdg, mat)

tests/test_rollout.py

@@ -9,23 +9,19 @@ import pytest
 import torch
 
 from giant.constants import TERM_ESCAPED, TERM_MAX_STEPS, TERM_UNKNOWN_PDG, K_MAX
+from giant.data.loader import TopNMap
 from giant.data.transforms import Normalizer
-from giant.model.network import Stage1Model, Stage2OneShot
-from giant.rollout import make_seed_frontier, rollout
+from giant.model.network import (
+    Stage1Model,
+    Stage2Autoregressive,
+    Stage2OneShot,
+    stage2_trunk_sec_dim,
+)
+from giant.rollout import L1DistCollector, make_seed_frontier, rollout
 
 pytest.importorskip("sklearn")
 from giant import geometry as g  # noqa: E402
 
-# giant/rollout.py isn't updated yet — it drives Stage1Model/Stage2OneShot
-# through giant.sample's sample_flow/sample_secondaries, which still call
-# models with the pre-refactor positional convention
-# (model(x, t, cond_cont, cond_cat)) that no longer matches these classes'
-# forward signatures. Deferred to docs/v0.3.0-design.md step 6/§10.
-pytestmark = pytest.mark.xfail(
-    reason="giant/rollout.py not updated for v0.3.0 network.py yet (step 6)",
-    strict=False,
-)
-
 PDG_MAP = {22: 0, 11: 1, -11: 2}
 MAT_MAP = {"G4_AIR": 0, "G4_PbWO4": 1}
 
@@ -350,3 +346,248 @@ def test_on_chunk_never_buffers_full_records():
     assert rec.termination_reason_counts == {"natural_end": 1}
     with pytest.raises(AssertionError):
         rec.to_dict()
+
+
+# ── v0.3.0 step 6: per-stage generators, AR decoder, particle_type.target ───
+
+# emb_dim=3: "other" (class idx 2) is shared by -11 and 13 (muon), matching
+# the real shape build_pdg_topn_map_from_files produces — see
+# decode_topn_class's docstring.
+PDG_TOPN_MAP = TopNMap(
+    class_map={22: 0, 11: 1, -11: 2, 13: 2},
+    other_members={-11: 5, 13: 1},
+)
+
+
+def _models_v3(
+    conditioning="physical",
+    decoder="one_shot",
+    target="physical",
+    generator1="flow",
+    generator2="flow",
+    k_max=6,
+    emb_dim=4,
+    stage2_has_n_sec_head=True,
+):
+    particle_cfg = {"type": conditioning, "emb_dim": emb_dim, "n_layers": 1}
+    material_cfg = {"type": conditioning, "emb_dim": emb_dim, "n_layers": 1}
+    # A fresh v0.3.0 Stage1Model — no n_sec_head_k_max, unlike _models() above
+    # (decision 1 moves n_sec ownership to stage 2 by default).
+    s1 = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+        generator=generator1,
+        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,
+    )
+    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)
+    else:
+        s2 = Stage2Autoregressive(**common)
+    return s1.eval(), s2.eval()
+
+
+def _run_v3(
+    s1,
+    s2,
+    escape_threshold=1e9,
+    energy_cutoff=1.0,
+    max_steps=15,
+    max_tracks_per_event=100,
+    seeds=None,
+    conditioning="physical",
+    pdg_topn_map=None,
+    other_policy="sample",
+    seed=0,
+    stage1_ddpm_steps=1000,
+    l1_dist_collector=None,
+):
+    torch.manual_seed(0)
+    np.random.seed(0)
+    cond, tgt, sec_phys = _norms()
+    return rollout(
+        s1,
+        s2,
+        _oracle(),
+        seeds or _seeds(),
+        cond,
+        tgt,
+        sec_phys,
+        PDG_MAP,
+        MAT_MAP,
+        energy_cutoff=energy_cutoff,
+        max_steps=max_steps,
+        steps=3,
+        batch_size=128,
+        max_tracks_per_event=max_tracks_per_event,
+        escape_threshold=escape_threshold,
+        conditioning=conditioning,
+        pdg_topn_map=pdg_topn_map,
+        other_policy=other_policy,
+        seed=seed,
+        stage1_ddpm_steps=stage1_ddpm_steps,
+        l1_dist_collector=l1_dist_collector,
+    )
+
+
+def test_rollout_stage2_owns_n_sec_when_stage1_has_no_head(fake_material_props):
+    """A fresh v0.3.0 Stage1Model has no n_sec_head (decision 1) — n_sec
+    must come from Stage2's own head instead, and the run must still
+    complete and conserve energy."""
+    s1, s2 = _models_v3()
+    rec = _run_v3(s1, s2)
+    assert len(rec["event_id"]) > 0
+    seeds = _seeds()
+    for i, ev in enumerate(seeds["event_id"]):
+        m = rec["event_id"] == ev
+        dep = rec["edep"][m].sum()
+        leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum()
+        assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4)
+
+
+def test_resolve_n_sec_raises_when_neither_stage_owns_head(fake_material_props):
+    """Neither stage owning n_sec_head only happens for a
+    stage2_model.n_sec.mode other than "head" — not a valid rollout-capable
+    checkpoint, and must fail with a clear error rather than crash deep
+    inside predict_n_sec."""
+    s1, s2 = _models_v3(stage2_has_n_sec_head=False)
+    with pytest.raises(RuntimeError, match="n_sec_head"):
+        _run_v3(s1, s2)
+
+
+@pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"])
+@pytest.mark.parametrize("generator2", ["flow", "wgan"])
+def test_rollout_physical_target_decoder_generator_matrix(
+    fake_material_props, decoder, generator2
+):
+    """Every (decoder, stage2 generator) combination under
+    particle_type.target="physical" must run to completion and conserve
+    energy — the matrix docs/v0.3.0-design.md §7 calls out for comparison."""
+    s1, s2 = _models_v3(decoder=decoder, generator2=generator2)
+    rec = _run_v3(s1, s2)
+    assert len(rec["event_id"]) > 0
+    seeds = _seeds()
+    for i, ev in enumerate(seeds["event_id"]):
+        m = rec["event_id"] == ev
+        dep = rec["edep"][m].sum()
+        leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum()
+        assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4)
+
+
+def test_sample_stage1_dispatches_ddpm_by_generator_kind():
+    """stage1_model.generator="ddpm" must be dispatched to sample_ddpm
+    (previously _step_chunk silently fell through to the flow ODE sampler
+    regardless of the checkpoint's actual generator — see giant.sample.sample_stage1).
+    A short T avoids the reverse-diffusion numerical blowup an untrained,
+    random-weight network produces over many steps; that instability is a
+    property of sampling from an untrained net, not of the dispatch logic
+    under test here, so a full oracle-driven rollout isn't needed."""
+    from giant.sample import sample_stage1
+
+    s1, _ = _models_v3(generator1="ddpm")
+    cond_cont = torch.randn(6, 15)
+    cond_cat = torch.zeros(6, 2, dtype=torch.long)
+    sample, n_sec = sample_stage1(s1, cond_cont, cond_cat, steps=10, ddpm_steps=5)
+    assert sample.shape == (6, 9)
+    assert n_sec is None  # fresh v0.3.0 Stage1Model owns no n_sec_head
+
+
+@pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"])
+def test_rollout_onehot_target_end_to_end(fake_material_props, decoder):
+    """particle_type.target="onehot" resolves a concrete PDG via
+    decode_topn_class (argmax + other_policy), and that PDG's real physics
+    (giant.particles.particle_phys_array) become the secondary's identity —
+    unlike "physical", not just a reporting label."""
+    s1, s2 = _models_v3(decoder=decoder, target="onehot", emb_dim=3)
+    rec = _run_v3(s1, s2, pdg_topn_map=PDG_TOPN_MAP, other_policy="modal")
+    assert len(rec["event_id"]) > 0
+    # Every spawned secondary's nominal pdg must be one decode_topn_class can
+    # actually produce (the topn map's known classes + its "other" members).
+    possible = set(PDG_TOPN_MAP.class_map.keys()) | set(
+        PDG_TOPN_MAP.other_members.keys()
+    )
+    secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())
+    assert secondary_pdgs <= possible
+
+
+def test_rollout_onehot_target_missing_topn_map_raises(fake_material_props):
+    s1, s2 = _models_v3(target="onehot", emb_dim=3)
+    with pytest.raises(RuntimeError, match="pdg_topn_map"):
+        _run_v3(s1, s2, pdg_topn_map=None)
+
+
+@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
+    conditioning's own particle embedding table, so every resolved PDG must
+    be a real member of the dense training vocab (pdg_map) — unlike
+    "onehot", there is no "other" bucket to fall outside of."""
+    s1, s2 = _models_v3(
+        conditioning="embedding", decoder=decoder, target="embedding", emb_dim=4
+    )
+    rec = _run_v3(s1, s2, conditioning="embedding")
+    assert len(rec["event_id"]) > 0
+    secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())
+    assert secondary_pdgs <= set(PDG_MAP.keys())
+
+
+def test_l1_dist_collector_populated_only_for_embedding_target():
+    """§11.3: the L1-distance diagnostic only makes sense under
+    particle_type.target="embedding" — a physical-target run must leave the
+    collector empty rather than silently accumulating garbage."""
+    s1, s2 = _models_v3(target="physical")
+    collector = L1DistCollector()
+    _run_v3(s1, s2, l1_dist_collector=collector)
+    assert collector.n == 0
+    assert collector.summary() is None
+
+
+def test_l1_dist_collector_accumulates_for_embedding_target():
+    s1, s2 = _models_v3(conditioning="embedding", target="embedding", emb_dim=4)
+    collector = L1DistCollector()
+    rec = _run_v3(s1, s2, conditioning="embedding", l1_dist_collector=collector)
+    n_secondaries = int((rec["generation"] > 0).sum())
+    assert n_secondaries > 0  # sanity: the tiny model does spawn secondaries
+    summary = collector.summary()
+    assert summary is not None
+    assert summary["n"] == collector.n > 0
+    assert summary["min"] <= summary["mean"] <= summary["max"]
+    assert summary["std"] >= 0.0
+    assert len(summary["hist_edges"]) == len(summary["hist_counts"]) + 1
+    assert sum(summary["hist_counts"]) <= summary["n"]  # some may fall outside [lo, hi)
+
+
+def test_l1_dist_collector_add_ignores_invalid_slots():
+    collector = L1DistCollector()
+    dist = np.array([[1.0, 5.0, 9.0]])
+    valid = np.array([[True, False, True]])
+    collector.add(dist, valid)
+    assert collector.n == 2
+    assert collector.minimum == 1.0
+    assert collector.maximum == 9.0
+
+
+def test_l1_dist_collector_add_empty_is_noop():
+    collector = L1DistCollector()
+    collector.add(np.zeros((0, 3)), np.zeros((0, 3), dtype=bool))
+    assert collector.n == 0
+    assert collector.summary() is None

tests/test_router.py

@@ -869,21 +869,17 @@ def test_build_models_routed_stage2_ties_to_stage1_router():
     assert stage1.trunk.router is stage2.trunk.router
 
 
-@pytest.mark.xfail(
-    reason=(
-        "giant/sample.py isn't updated yet — its sample_flow/sample_secondaries "
-        "call models positionally as model(x, t, cond_cont, cond_cat), which "
-        "doesn't match Stage1Model/Stage2OneShot's new forward signature. "
-        "Deferred to docs/v0.3.0-design.md step 6."
-    ),
-    strict=False,
-)
 def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow():
     from giant.sample import sample_flow, sample_secondaries
 
     cfg = _nested_cfg(
         pdg_vocab=3,
         mat_vocab=2,
+        # PdgRouter (axis1) always builds its own training-vocab embedding,
+        # incompatible with conditioning.particle.type="physical" (the
+        # _nested_cfg default) — see _check_router_conditioning_compat.
+        particle_type="embedding",
+        material_type="embedding",
         stage1_router={
             "enabled": True,
             "type": "composed",
@@ -900,6 +896,11 @@ def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow():
     cond_cont, cond_cat = _cond(B, pdg=3, mat=2)
     stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2)
     assert stage1_norm.shape == (B, X_DIM)
+    # A fresh v0.3.0 Stage1Model has no n_sec_head (decision 1 moves it to
+    # stage 2) — sample_flow returns n_sec_pred=None here, and n_sec must be
+    # asked of stage2 instead, using the just-sampled stage1_norm as context.
+    assert n_sec_pred is None
+    n_sec_pred = stage2.predict_n_sec(cond_cont, cond_cat, stage1_norm).argmax(dim=-1)
     assert n_sec_pred.shape == (B,)
 
     sec_cont, sec_type_emb, sec_valid = sample_secondaries(
@@ -1075,15 +1076,6 @@ def test_build_models_routed_when_enabled():
     assert len(stage2.trunk.experts) == 4
 
 
-@pytest.mark.xfail(
-    reason=(
-        "giant/sample.py isn't updated yet — its sample_flow/sample_secondaries "
-        "call models positionally as model(x, t, cond_cont, cond_cat), which "
-        "doesn't match Stage1Model/Stage2OneShot's new forward signature. "
-        "Deferred to docs/v0.3.0-design.md step 6."
-    ),
-    strict=False,
-)
 def test_build_models_routed_pair_is_drop_in_for_sample_flow():
     """Exercise the exact calling convention giant/sample.py uses."""
     from giant.sample import sample_flow, sample_secondaries
@@ -1100,6 +1092,11 @@ def test_build_models_routed_pair_is_drop_in_for_sample_flow():
     cond_cont, cond_cat = _cond(B, pdg=3, mat=2)
     stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2)
     assert stage1_norm.shape == (B, X_DIM)
+    # A fresh v0.3.0 Stage1Model has no n_sec_head (decision 1 moves it to
+    # stage 2) — sample_flow returns n_sec_pred=None here, and n_sec must be
+    # asked of stage2 instead, using the just-sampled stage1_norm as context.
+    assert n_sec_pred is None
+    n_sec_pred = stage2.predict_n_sec(cond_cont, cond_cat, stage1_norm).argmax(dim=-1)
     assert n_sec_pred.shape == (B,)
 
     sec_cont, sec_type_emb, sec_valid = sample_secondaries(

tests/test_sample.py

@@ -0,0 +1,237 @@
+"""Tests for giant/sample.py's v0.3.0 stage-model sampling — the AR loop
+(`sample_secondaries_ar`) and non-"physical" `particle_type.target` coverage
+for the one-shot samplers (docs/v0.3.0-design.md step 6)."""
+
+import pytest
+import torch
+
+from giant.constants import COND_DIM, CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, X_DIM
+from giant.model.network import (
+    Stage1Model,
+    Stage2Autoregressive,
+    Stage2OneShot,
+    stage2_trunk_sec_dim,
+)
+from giant.sample import (
+    sample_flow,
+    sample_secondaries,
+    sample_secondaries_ar,
+    sample_secondaries_wgan,
+    sample_wgan,
+)
+
+_PHYS_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
+
+
+def _particle_material_cfg(conditioning: str, emb_dim: int) -> tuple[dict, dict]:
+    cfg = {"type": conditioning, "emb_dim": emb_dim, "n_layers": 1}
+    return dict(cfg), dict(cfg)
+
+
+def _cond(B: int, pdg: int = 3, mat: int = 2) -> tuple[torch.Tensor, torch.Tensor]:
+    cond_cont = torch.randn(B, COND_DIM)
+    cond_cat = torch.stack(
+        [torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
+    )
+    return cond_cont, cond_cat
+
+
+def _conditioning_for(target: str) -> str:
+    # target="embedding" regresses against the conditioning's own embedding
+    # table (docs/v0.3.0-design.md §3.3) — only meaningful when the
+    # conditioning axis is itself "embedding".
+    return "embedding" if target == "embedding" else "physical"
+
+
+def _stage2_oneshot(
+    target: str, generator: str, emb_dim: int = 6, pdg: int = 3, mat: int = 2
+) -> Stage2OneShot:
+    particle_cfg, material_cfg = _particle_material_cfg(
+        _conditioning_for(target), emb_dim
+    )
+    particle_type_cfg = {"target": target}
+    # build_models (giant/model/network.py) computes sec_dim this same way
+    # before constructing Stage2OneShot — its own default (SEC_DIM, the
+    # "physical" width) is only correct for target="physical".
+    sec_dim = stage2_trunk_sec_dim(particle_type_cfg, generator, K_MAX, emb_dim)
+    return Stage2OneShot(
+        pdg_vocab=pdg,
+        mat_vocab=mat,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+        generator=generator,
+        time_dim=16,
+        noise_dim=8,
+        sec_dim=sec_dim,
+        particle_type_cfg=particle_type_cfg,
+    ).eval()
+
+
+def _stage2_ar(
+    target: str,
+    generator: str,
+    emb_dim: int = 6,
+    pdg: int = 3,
+    mat: int = 2,
+    k_max: int = 5,
+) -> Stage2Autoregressive:
+    particle_cfg, material_cfg = _particle_material_cfg(
+        _conditioning_for(target), emb_dim
+    )
+    return Stage2Autoregressive(
+        pdg_vocab=pdg,
+        mat_vocab=mat,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+        generator=generator,
+        time_dim=16,
+        noise_dim=8,
+        k_max=k_max,
+        particle_type_cfg={"target": target},
+    ).eval()
+
+
+def _expected_type_dim(target: str, emb_dim: int) -> int:
+    return PARTICLE_PHYS_DIM if target == "physical" else emb_dim
+
+
+# ── Stage-1 n_sec ownership (decision 1) ────────────────────────────────────
+
+
+def test_sample_flow_returns_none_n_sec_when_stage1_owns_no_head():
+    model = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=_PHYS_CFG,
+        material_cfg=_PHYS_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+    )
+    cond_cont, cond_cat = _cond(4)
+    sample, n_sec = sample_flow(model, cond_cont, cond_cat, steps=2)
+    assert sample.shape == (4, X_DIM)
+    assert n_sec is None
+
+
+def test_sample_wgan_returns_none_n_sec_when_stage1_owns_no_head():
+    model = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=_PHYS_CFG,
+        material_cfg=_PHYS_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        generator="wgan",
+        noise_dim=8,
+    )
+    cond_cont, cond_cat = _cond(4)
+    sample, n_sec = sample_wgan(model, cond_cont, cond_cat)
+    assert sample.shape == (4, X_DIM)
+    assert n_sec is None
+
+
+def test_sample_flow_returns_n_sec_for_legacy_stage1():
+    model = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=_PHYS_CFG,
+        material_cfg=_PHYS_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        n_sec_head_k_max=K_MAX,
+    )
+    cond_cont, cond_cat = _cond(5)
+    _, n_sec = sample_flow(model, cond_cont, cond_cat, steps=2)
+    assert n_sec.shape == (5,)
+
+
+# ── Stage2OneShot: non-"physical" particle_type.target ──────────────────────
+
+
+@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
+def test_sample_secondaries_flow_shapes_by_target(target):
+    B, emb_dim = 5, 6
+    decoder = _stage2_oneshot(target, "flow", emb_dim=emb_dim)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
+    sec_cont, sec_type, sec_valid = sample_secondaries(
+        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
+    )
+    assert sec_cont.shape == (B, K_MAX, CONT_SLOT_DIM)
+    assert sec_type.shape == (B, K_MAX, _expected_type_dim(target, emb_dim))
+    assert sec_valid.shape == (B, K_MAX)
+    assert torch.isfinite(sec_cont).all()
+    assert torch.isfinite(sec_type).all()
+
+
+@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
+def test_sample_secondaries_wgan_shapes_by_target(target):
+    B, emb_dim = 5, 6
+    decoder = _stage2_oneshot(target, "wgan", emb_dim=emb_dim)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
+    sec_cont, sec_type, sec_valid = sample_secondaries_wgan(
+        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred
+    )
+    assert sec_cont.shape == (B, K_MAX, CONT_SLOT_DIM)
+    assert sec_type.shape == (B, K_MAX, _expected_type_dim(target, emb_dim))
+    assert sec_valid.shape == (B, K_MAX)
+
+
+# ── Stage2Autoregressive ─────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("generator", ["flow", "wgan"])
+@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
+def test_sample_secondaries_ar_shapes(target, generator):
+    B, k_max, emb_dim = 4, 5, 6
+    decoder = _stage2_ar(target, generator, emb_dim=emb_dim, k_max=k_max)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec_pred = torch.randint(0, k_max + 1, (B,))
+    sec_cont, sec_type, sec_valid = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
+    )
+    assert sec_cont.shape == (B, k_max, CONT_SLOT_DIM)
+    assert sec_type.shape == (B, k_max, _expected_type_dim(target, emb_dim))
+    assert sec_valid.shape == (B, k_max)
+    assert torch.isfinite(sec_cont).all()
+    assert torch.isfinite(sec_type).all()
+
+
+@pytest.mark.parametrize("generator", ["flow", "wgan"])
+@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
+def test_sample_secondaries_ar_valid_mask_matches_n_sec(target, generator):
+    B, k_max, emb_dim = 3, 5, 6
+    decoder = _stage2_ar(target, generator, emb_dim=emb_dim, k_max=k_max)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec_pred = torch.tensor([0, 2, k_max])
+    _, _, sec_valid = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
+    )
+    for i, n in enumerate(n_sec_pred.tolist()):
+        assert sec_valid[i, :n].all()
+        assert not sec_valid[i, n:].any()
+
+
+def test_sample_secondaries_ar_first_slot_has_no_history():
+    """Slot 0 always has has_prev=False internally — nothing to assert on
+    the public API directly, but a k_max=1 run should not crash on the
+    "previous token" path at all (has_prev never true)."""
+    B, emb_dim = 3, 6
+    decoder = _stage2_ar("physical", "flow", emb_dim=emb_dim, k_max=1)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec_pred = torch.tensor([0, 1, 1])
+    sec_cont, sec_type, sec_valid = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
+    )
+    assert sec_cont.shape == (B, 1, CONT_SLOT_DIM)
+    assert sec_valid.tolist() == [[False], [True], [True]]

tests/test_type_embedding_distance.py

@@ -0,0 +1,43 @@
+"""Tests for the secondary-type embedding-distance diagnostic
+(giant.analysis.type_embedding_distance) — the §11.3 diagnostic."""
+
+from __future__ import annotations
+
+from giant.analysis.type_embedding_distance import compute_type_embedding_l1_distance
+
+
+def _summary(n=100):
+    return {
+        "n": n,
+        "mean": 1.23,
+        "std": 0.45,
+        "min": 0.01,
+        "max": 9.87,
+        "hist_edges": [0.0, 1.0, 2.0, 3.0],
+        "hist_counts": [30, 40, 30],
+    }
+
+
+def test_none_is_unavailable():
+    r = compute_type_embedding_l1_distance(None)
+    assert r.kind == "unavailable"
+    assert r.id == "type_embedding_l1_distance"
+    assert r.payload["note"]
+
+
+def test_summary_produces_single_hist():
+    r = compute_type_embedding_l1_distance(_summary())
+    assert r.kind == "single_hist"
+    assert r.id == "type_embedding_l1_distance"
+    assert r.payload["edges"] == [0.0, 1.0, 2.0, 3.0]
+    assert r.payload["rollout"] == [30, 40, 30]
+    assert r.payload["log_x"] is True
+    assert r.payload["log_y"] is True
+    assert "n=100" in r.payload["note"]
+
+
+def test_single_hist_payload_shape_matches_render_contract():
+    """_render_single (giant.analysis.render) requires len(rollout) ==
+    len(edges) - 1."""
+    r = compute_type_embedding_l1_distance(_summary())
+    assert len(r.payload["rollout"]) == len(r.payload["edges"]) - 1

tests/test_validate.py

@@ -1,5 +1,4 @@
 import numpy as np
-import pytest
 import torch
 
 from giant.constants import COND_DIM, K_MAX, SEC_SLOT_DIM, X_DIM
@@ -49,15 +48,6 @@ def _zero_secondaries_loader(B=4, n_batches=2):
     return batches
 
 
-@pytest.mark.xfail(
-    reason=(
-        "giant/validate.py isn't updated yet — it calls the stage models "
-        "(sample_secondaries et al.) with the old positional convention, "
-        "which doesn't match Stage1Model/Stage2OneShot's new forward "
-        "signature. Deferred to docs/v0.3.0-design.md step 6/§10."
-    ),
-    strict=False,
-)
 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