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>

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