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

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

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

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

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

tests/test_validate.py

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