Decouple secondary-species vocabulary from conditioning.particle.emb_dim (gitea #29)

conditioning.particle.emb_dim and stage2_model.particle_type.target="onehot"'s
class count were silently the same number everywhere (pipeline.py's PDG
top-N map build, Stage2OneShot/Stage2Autoregressive's type head, StageSpec's
training loss width, the checkpoint's shared pdg_topn_map), fixing the
secondary-species vocabulary at whatever width the unrelated
physical-conditioning MLP happened to use — the exact vocabulary the v0.3.0
pivot exists to fix.

Adds stage2_model.particle_type.n_classes (default 0 = inherit
conditioning.particle.emb_dim, preserving today's behavior and every
existing checkpoint) and a single resolve_type_n_classes helper used
everywhere the coupling used to be implicit. Splits the checkpoint's shared
pdg_topn_map into a conditioning-only pdg_topn_map and a new
sec_type_topn_map, built independently through the existing
(axis, n_classes)-keyed setup cache (no extra scan when they still resolve
to the same N) and threaded through giant predict/giant rollout's decode
path. A checkpoint with no sec_type_topn_map key (pre-#29) falls back to
reusing pdg_topn_map, reproducing the old shared behavior exactly.

Decided with the user during planning: commit directly on this branch;
represent the split as an additive sec_type_topn_map checkpoint key rather
than conditionally reusing pdg_topn_map; build the two top-N maps
independently rather than the issue's proposed build-at-max-and-slice, since
the setup cache already avoids redundant scans across runs.

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

tests/test_network.py

@@ -288,6 +288,33 @@ def test_stage2_oneshot_forward_shape_onehot_flow_excludes_type():
     assert out.shape == (B, k_max * CONT_SLOT_DIM)
 
 
+def test_stage2_oneshot_particle_type_n_classes_overrides_conditioning_emb_dim():
+    """gitea #29: stage2_model.particle_type.n_classes, not
+    conditioning.particle.emb_dim, sizes the onehot type_head/type_dim when
+    explicitly set — the two used to be silently the same number."""
+    k_max = 5
+    particle_cfg = {"type": "physical", "emb_dim": 6, "n_layers": 1}
+    particle_type_cfg = {"target": "onehot", "lambda": 1.0, "n_classes": 20}
+    sec_dim = stage2_trunk_sec_dim(particle_type_cfg, "flow", k_max, 20)
+    model = Stage2OneShot(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=particle_cfg,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        cond_out_dim=16,
+        context_dim=8,
+        sec_dim=sec_dim,
+        generator="flow",
+        k_max=k_max,
+        particle_type_cfg=particle_type_cfg,
+    )
+    assert model.type_dim == 20  # not particle_cfg["emb_dim"] == 6
+    assert model.type_head is not None
+    assert model.type_head[-1].out_features == k_max * 20
+
+
 # --- MarkovHistory -----------------------------------------------------------
 
 
@@ -423,6 +450,29 @@ def test_stage2_autoregressive_history_invalid_raises():
         _build_stage2_ar("onehot", "wgan", history="bogus")
 
 
+def test_stage2_autoregressive_particle_type_n_classes_overrides_conditioning_emb_dim():
+    """gitea #29, Stage2Autoregressive side — see the Stage2OneShot version
+    of this test for the full rationale."""
+    particle_cfg = {"type": "physical", "emb_dim": 6, "n_layers": 1}
+    particle_type_cfg = {"target": "onehot", "lambda": 1.0, "n_classes": 20}
+    model = Stage2Autoregressive(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=particle_cfg,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        cond_out_dim=16,
+        context_dim=8,
+        generator="flow",
+        k_max=5,
+        particle_type_cfg=particle_type_cfg,
+    )
+    assert model.type_dim == 20  # not particle_cfg["emb_dim"] == 6
+    assert model.type_head is not None
+    assert model.type_head[-1].out_features == 20
+
+
 @pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
 @pytest.mark.parametrize("generator", ["wgan", "flow"])
 @pytest.mark.parametrize("history", ["markov", "attention"])
@@ -655,6 +705,33 @@ def test_build_models_share_stages_true_shared_params_are_in_both_stage_paramete
     assert shared_ids <= {id(p) for p in stage2.parameters()}
 
 
+def test_build_models_particle_type_n_classes_overrides_conditioning_emb_dim():
+    """gitea #29 end-to-end through build_models: setting
+    stage2_model.particle_type.n_classes independently of
+    conditioning.particle.emb_dim actually resizes the built stage2 model,
+    not just the two lower-level unit tests above."""
+    cfg = _minimal_model_config(share_stages=False)  # conditioning.particle.emb_dim = 4
+    cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0, "n_classes": 11}
+    built = build_models(cfg)
+    assert built["stage2"] is not None
+    assert built["stage2"].type_dim == 11
+
+
+def test_build_critics_particle_type_n_classes_overrides_conditioning_emb_dim():
+    cfg = _minimal_model_config(share_stages=False)  # conditioning.particle.emb_dim = 4
+    cfg["stage2_model"]["generator"] = "wgan"
+    cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0, "n_classes": 4}
+    default_n_classes_critic = build_critics(cfg)["stage2"]
+    assert default_n_classes_critic is not None
+
+    cfg["stage2_model"]["particle_type"]["n_classes"] = 11
+    wider_critic = build_critics(cfg)["stage2"]
+    assert wider_critic is not None
+    # k_max=3 slots, each CONT_SLOT_DIM + n_classes wide under wgan folding —
+    # widening n_classes alone (emb_dim stays 4) must widen the critic input.
+    assert wider_critic.input_proj.in_features > default_n_classes_critic.input_proj.in_features
+
+
 # ── build_models/build_critics: DEFAULT_CONFIG fallback drift (issues.md #1) ─