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_rollout.py

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