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>

giant/checkpoint_io.py

@@ -88,6 +88,23 @@ def load_mat_topn_map(ckpt: dict) -> TopNMap | None:
     return topnmap_from_json(raw, axis="material") if raw is not None else None
 
 
+def load_sec_type_topn_map(ckpt: dict) -> TopNMap | None:
+    """`ckpt["sec_type_topn_map"]` as a `giant.data.loader.TopNMap`, or
+    `None` if this checkpoint's `stage2_model.particle_type.target` was never
+    `"onehot"` (see `giant.pipeline.run_setup_stage`).
+
+    Pre-gitea-#29 checkpoints have no `sec_type_topn_map` key at all — before
+    #29, the secondary-species decode map and the conditioning PDG onehot map
+    were always numerically the same map, saved once under `pdg_topn_map`.
+    For those, fall back to `load_pdg_topn_map` to reproduce that exact
+    behavior; a current checkpoint always has the key (possibly `null`, if
+    `particle_type.target != "onehot"`), so this fallback never fires for one."""
+    if "sec_type_topn_map" in ckpt:
+        raw = ckpt["sec_type_topn_map"]
+        return topnmap_from_json(raw, axis="pdg") if raw is not None else None
+    return load_pdg_topn_map(ckpt)
+
+
 @dataclass(frozen=True)
 class InferenceContext:
     """Everything needed to run a trained checkpoint forward, resolved once."""
@@ -101,6 +118,7 @@ class InferenceContext:
     mat_map: dict[str, int]
     pdg_topn_map: TopNMap | None
     mat_topn_map: TopNMap | None
+    sec_type_topn_map: TopNMap | None
     particle_conditioning: str
     material_conditioning: str
     k_max: int
@@ -153,6 +171,13 @@ def load_for_inference(
         raise CheckpointCompatibilityError(
             "checkpoint's conditioning.material.type='onehot' but has no mat_topn_map — retrain with the current code"
         )
+    sec_type_topn_map = load_sec_type_topn_map(ckpt)
+    particle_type_target = stage_cfg(model_cfg, "stage2").get("particle_type", {}).get("target", "onehot")
+    if particle_type_target == "onehot" and sec_type_topn_map is None:
+        raise CheckpointCompatibilityError(
+            "checkpoint's stage2_model.particle_type.target='onehot' but has no "
+            "sec_type_topn_map — retrain with the current code"
+        )
     other_policy = particle_type_other_policy(model_cfg)
     stage1_ddpm_steps = ddpm_steps(model_cfg, "stage1")
     stage2_ddpm_steps = ddpm_steps(model_cfg, "stage2")
@@ -196,6 +221,7 @@ def load_for_inference(
         mat_map=mat_map,
         pdg_topn_map=pdg_topn_map,
         mat_topn_map=mat_topn_map,
+        sec_type_topn_map=sec_type_topn_map,
         particle_conditioning=particle_conditioning,
         material_conditioning=material_conditioning,
         k_max=k_max,

giant/cli.py

@@ -953,6 +953,7 @@ def predict(
     cond_norm, tgt_norm, sec_phys_norm = ctx.cond_norm, ctx.tgt_norm, ctx.sec_phys_norm
     pdg_map, mat_map = ctx.pdg_map, ctx.mat_map
     pdg_topn_map, mat_topn_map = ctx.pdg_topn_map, ctx.mat_topn_map
+    sec_type_topn_map = ctx.sec_type_topn_map
     particle_conditioning, material_conditioning = ctx.particle_conditioning, ctx.material_conditioning
     other_policy = ctx.other_policy
     stage1_ddpm_steps = ctx.stage1_ddpm_steps
@@ -1105,7 +1106,7 @@ def predict(
                 piece["pre_dir"],
                 sec_phys_norm,
                 pdg_map,
-                pdg_topn_map,
+                sec_type_topn_map,
                 other_policy,
                 None,
             )
@@ -1326,6 +1327,7 @@ def rollout(
     cond_norm, tgt_norm, sec_phys_norm = ctx.cond_norm, ctx.tgt_norm, ctx.sec_phys_norm
     pdg_map, mat_map = ctx.pdg_map, ctx.mat_map
     pdg_topn_map, mat_topn_map = ctx.pdg_topn_map, ctx.mat_topn_map
+    sec_type_topn_map = ctx.sec_type_topn_map
     particle_conditioning, material_conditioning = ctx.particle_conditioning, ctx.material_conditioning
     other_policy = ctx.other_policy
     stage1_ddpm_steps, stage2_ddpm_steps = ctx.stage1_ddpm_steps, ctx.stage2_ddpm_steps
@@ -1386,6 +1388,7 @@ def rollout(
         material_conditioning=material_conditioning,
         pdg_topn_map=pdg_topn_map,
         mat_topn_map=mat_topn_map,
+        sec_type_topn_map=sec_type_topn_map,
         other_policy=other_policy,
         seed=seed,
         stage1_ddpm_steps=stage1_ddpm_steps,

giant/config.py

@@ -418,6 +418,11 @@ class ParticleTypeConfig:
     # at map-build time. "modal": always the most common member. "drop":
     # discard the secondary. Read only under target = "onehot".
     other_policy: str = "sample"
+    # Secondary-species class count under target = "onehot" — independent of
+    # conditioning.particle.emb_dim (see gitea #29: the two used to be
+    # silently the same number). 0 = inherit conditioning.particle.emb_dim,
+    # preserving pre-#29 behavior.
+    n_classes: int = 0
 
     @classmethod
     def from_dict(cls, d: dict | None) -> "ParticleTypeConfig":
@@ -426,10 +431,16 @@ class ParticleTypeConfig:
             target=d.get("target", "onehot"),
             lambda_weight=d.get("lambda", 1.0),
             other_policy=d.get("other_policy", "sample"),
+            n_classes=d.get("n_classes", 0),
         )
 
     def to_dict(self) -> dict:
-        return {"target": self.target, "lambda": self.lambda_weight, "other_policy": self.other_policy}
+        return {
+            "target": self.target,
+            "lambda": self.lambda_weight,
+            "other_policy": self.other_policy,
+            "n_classes": self.n_classes,
+        }
 
 
 @dataclass(frozen=True)

giant/model/builders.py

@@ -7,7 +7,14 @@ from giant.config import ConditioningConfig, Stage1ModelConfig, Stage2ModelConfi
 from giant.constants import X_DIM
 from giant.model._legacy import _migrate_legacy_model_config
 from giant.model.encoders import ConditionEncoder
-from giant.model.models import CriticModel, Stage1Model, Stage2Autoregressive, Stage2OneShot, stage2_trunk_sec_dim
+from giant.model.models import (
+    CriticModel,
+    Stage1Model,
+    Stage2Autoregressive,
+    Stage2OneShot,
+    resolve_type_n_classes,
+    stage2_trunk_sec_dim,
+)
 from giant.model.routers import Router, _build_router_from_cfg
 
 # ---------------------------------------------------------------------------
@@ -122,7 +129,9 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 cond_enc=shared_cond_enc,
             )
         else:
-            sec_dim = stage2_trunk_sec_dim(particle_type_cfg, generator, k_max, particle_cfg["emb_dim"])
+            sec_dim = stage2_trunk_sec_dim(
+                particle_type_cfg, generator, k_max, resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
+            )
             result["stage2"] = Stage2OneShot(
                 pdg_vocab=pdg_vocab,
                 mat_vocab=mat_vocab,
@@ -182,7 +191,9 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
     if s2_spec.active and s2_spec.generator == "wgan":
         k_max = s2_spec.k_max
         particle_type_cfg = s2_spec.particle_type.to_dict()
-        in_dim = stage2_trunk_sec_dim(particle_type_cfg, "wgan", k_max, particle_cfg["emb_dim"])
+        in_dim = stage2_trunk_sec_dim(
+            particle_type_cfg, "wgan", k_max, resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
+        )
         result["stage2"] = CriticModel(
             pdg_vocab=pdg_vocab,
             mat_vocab=mat_vocab,

giant/model/models.py

@@ -16,11 +16,28 @@ from giant.model.trunks import build_trunk
 # ---------------------------------------------------------------------------
 
 
+def resolve_type_n_classes(particle_type_cfg: dict, particle_emb_dim: int) -> int:
+    """Effective width fed to `stage2_type_dim`/`stage2_trunk_sec_dim` in
+    place of a bare `conditioning.particle.emb_dim` read. Under
+    `target = "onehot"` this is `stage2_model.particle_type.n_classes` (0 =
+    inherit `conditioning.particle.emb_dim`) — see gitea #29, which decoupled
+    the secondary-species vocabulary size from the unrelated
+    physical-conditioning MLP's output width. Under `target = "embedding"`
+    (or `"physical"`, which ignores this value entirely) `n_classes` doesn't
+    apply — the width stays `conditioning.particle.emb_dim`, the embedding
+    table's own dimensionality (`validate_config` requires
+    `conditioning.particle.type = "embedding"` here)."""
+    if particle_type_cfg.get("target", "physical") == "onehot":
+        return particle_type_cfg.get("n_classes", 0) or particle_emb_dim
+    return particle_emb_dim
+
+
 def stage2_type_dim(particle_type_cfg: dict, emb_dim: int) -> int:
     """Width of a single secondary slot's type slice —
     `PARTICLE_PHYS_DIM` (log_mass, charge) for `target = "physical"`, else
     `emb_dim` (both `"onehot"` class logits and `"embedding"` vectors are
-    `conditioning.particle.emb_dim` wide)."""
+    this many classes/dims wide — callers resolve `emb_dim` via
+    `resolve_type_n_classes` first)."""
     target = particle_type_cfg.get("target", "physical")
     return PARTICLE_PHYS_DIM if target == "physical" else emb_dim
 
@@ -177,7 +194,9 @@ class Stage2OneShot(nn.Module):
         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.type_dim = stage2_type_dim(
+            self.particle_type_cfg, resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"])
+        )
         self.cond_enc = (
             cond_enc
             if cond_enc is not None
@@ -203,7 +222,7 @@ class Stage2OneShot(nn.Module):
         self.type_head = None
         target = self.particle_type_cfg.get("target", "physical")
         if target != "physical" and generator in ("flow", "ddpm"):
-            emb_dim = particle_cfg["emb_dim"]
+            emb_dim = resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"])
             self.type_head = nn.Sequential(
                 nn.Linear(cond_out_dim, hidden_dim // 2),
                 nn.SiLU(),
@@ -324,7 +343,7 @@ class Stage2Autoregressive(nn.Module):
         self.noise_dim = noise_dim
         self.k_max = k_max
         self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"})
-        emb_dim = particle_cfg["emb_dim"]
+        emb_dim = resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"])
         self.type_dim = stage2_type_dim(self.particle_type_cfg, emb_dim)
 
         self.cond_enc = (

giant/model/network.py

@@ -17,6 +17,7 @@ from giant.model.models import (
     Stage1Model,
     Stage2Autoregressive,
     Stage2OneShot,
+    resolve_type_n_classes,
     stage2_trunk_sec_dim,
     stage2_type_dim,
 )
@@ -80,6 +81,7 @@ __all__ = [
     "cat_col_layout",
     "migrate_legacy_state_dict",
     "register_router",
+    "resolve_type_n_classes",
     "stage2_trunk_sec_dim",
     "stage2_type_dim",
 ]

giant/particles.py

@@ -140,9 +140,8 @@ def decode_topn_class(
     other_policy: str = "sample",
     rng: np.random.Generator | None = None,
 ) -> np.ndarray:
-    """`conditioning.particle.type` / `stage2_model.particle_type.target =
-    "onehot"` inference decode: per-row top-N class index -> concrete PDG
-    code.
+    """`stage2_model.particle_type.target = "onehot"` inference decode:
+    per-row top-N class index -> concrete secondary-species 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`)
@@ -150,8 +149,11 @@ def decode_topn_class(
         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
+    n_classes: the resolved secondary-species class count
+        (`giant.model.models.resolve_type_n_classes` —
+        `stage2_model.particle_type.n_classes`, 0 = inherit
+        `conditioning.particle.emb_dim`; see gitea #29); 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"`

giant/pipeline.py

@@ -31,7 +31,7 @@ from giant.data.transforms import (
     sorted_membership,
 )
 from giant.data.dataset import make_event_split, StreamingStepsDataset
-from giant.model.network import build_models, build_critics
+from giant.model.network import build_models, build_critics, resolve_type_n_classes
 from giant.training import train as run_training
 
 
@@ -49,6 +49,7 @@ class SetupStageResult:
     mat_map: dict[str, int]
     proc_map: dict[str, int] | None
     pdg_topn_map: TopNMap | None
+    sec_type_topn_map: TopNMap | None
     mat_topn_map: TopNMap | None
     cond_norm: Normalizer
     tgt_norm: Normalizer
@@ -175,29 +176,42 @@ def run_setup_stage(
                 cache.proc_maps[n_experts] = proc_map
 
     # Top-N-plus-other maps for onehot conditioning/type axes.
-    # The PDG axis is shared by
-    # conditioning.particle.type="onehot" and
-    # stage2_model.particle_type.target="onehot" (both key off
-    # conditioning.particle.emb_dim), so at most one PDG scan is needed even
-    # if both consumers are active. The material axis is independent.
+    # The PDG axis is used independently by conditioning.particle.type="onehot"
+    # (cond_cat's onehot feature) and stage2_model.particle_type.target="onehot"
+    # (secondary-species decode) — their class counts can now differ (gitea
+    # #29: stage2_model.particle_type.n_classes, 0 = inherit
+    # conditioning.particle.emb_dim), so each is resolved and built
+    # independently via _pdg_topn below. cache.topn_maps is keyed by
+    # (axis, n_classes) (setup_cache.topn_key), so when the two resolve to
+    # the same N the second call is a cache hit against the first — no extra
+    # scan in the common case where they still match. The material axis is
+    # independent of both.
     particle_cfg = cfg["conditioning"]["particle"]
     material_cfg = cfg["conditioning"]["material"]
-    particle_type_target = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type")).target
+    particle_type_cfg_dict = cfg["stage2_model"].get("particle_type") or {}
+    particle_type_target = config.ParticleTypeConfig.from_dict(particle_type_cfg_dict).target
 
-    pdg_topn_map: TopNMap | None = None
-    if particle_cfg["type"] == "onehot" or particle_type_target == "onehot":
-        n_classes = particle_cfg["emb_dim"]
+    def _pdg_topn(n_classes: int) -> TopNMap:
         cache_key = setup_cache.topn_key("pdg", n_classes)
         cached = cache.topn_maps.get(cache_key) if cache is not None else None
         if cached is not None:
-            pdg_topn_map = cached
-            echo(f"pdg top-N map: cache hit ({len(pdg_topn_map.class_map)} codes, {n_classes} classes)")
-        else:
-            echo("building pdg top-N map …")
-            pdg_topn_map = build_pdg_topn_map_from_files(files, n_classes=n_classes)
-            echo(f"  {len(pdg_topn_map.class_map)} pdg codes mapped to {n_classes} classes")
-            if cache is not None:
-                cache.topn_maps[cache_key] = pdg_topn_map
+            echo(f"pdg top-N map: cache hit ({len(cached.class_map)} codes, {n_classes} classes)")
+            return cached
+        echo("building pdg top-N map …")
+        topn_map = build_pdg_topn_map_from_files(files, n_classes=n_classes)
+        echo(f"  {len(topn_map.class_map)} pdg codes mapped to {n_classes} classes")
+        if cache is not None:
+            cache.topn_maps[cache_key] = topn_map
+        return topn_map
+
+    pdg_topn_map: TopNMap | None = None
+    if particle_cfg["type"] == "onehot":
+        pdg_topn_map = _pdg_topn(particle_cfg["emb_dim"])
+
+    sec_type_topn_map: TopNMap | None = None
+    if particle_type_target == "onehot":
+        sec_type_n_classes = resolve_type_n_classes(particle_type_cfg_dict, particle_cfg["emb_dim"])
+        sec_type_topn_map = _pdg_topn(sec_type_n_classes)
 
     mat_topn_map: TopNMap | None = None
     if material_cfg["type"] == "onehot":
@@ -296,6 +310,7 @@ def run_setup_stage(
         mat_map=mat_map,
         proc_map=proc_map,
         pdg_topn_map=pdg_topn_map,
+        sec_type_topn_map=sec_type_topn_map,
         mat_topn_map=mat_topn_map,
         cond_norm=cond_norm,
         tgt_norm=tgt_norm,
@@ -385,8 +400,8 @@ def run_train_job(
     # (physical stays untouched/None).
     particle_type_target = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type")).target
     if particle_type_target == "onehot":
-        assert setup.pdg_topn_map is not None
-        sec_type_class_map = setup.pdg_topn_map.class_map
+        assert setup.sec_type_topn_map is not None
+        sec_type_class_map = setup.sec_type_topn_map.class_map
     elif particle_type_target == "embedding":
         sec_type_class_map = pdg_map
     else:
@@ -490,6 +505,7 @@ def run_train_job(
         mat_map={str(k): v for k, v in mat_map.items()},
         proc_map=proc_map,
         pdg_topn_map=setup.pdg_topn_map,
+        sec_type_topn_map=setup.sec_type_topn_map,
         mat_topn_map=setup.mat_topn_map,
         model_config=model_config,
         resume_path=resume,

giant/rollout.py

@@ -117,7 +117,7 @@ def decode_secondary_identity(
     pre_dir: np.ndarray,
     sec_phys_norm: Normalizer,
     pdg_map: dict[int, int],
-    pdg_topn_map: "TopNMap | None",
+    sec_type_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]:
@@ -158,15 +158,15 @@ def decode_secondary_identity(
     l1_dist = None
 
     if target == "onehot":
-        if pdg_topn_map is None:
+        if sec_type_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']"
+                "particle_type.target='onehot' rollout needs sec_type_topn_map "
+                "(the checkpoint's saved top-N map) — see ckpt['sec_type_topn_map']"
             )
         class_idx = sec_type_np.argmax(axis=-1)
         sec_pdg = decode_topn_class(
             class_idx,
-            pdg_topn_map,
+            sec_type_topn_map,
             n_classes=sec_decoder.type_dim,
             other_policy=other_policy,
             rng=rng,
@@ -444,6 +444,7 @@ def rollout(
     material_conditioning: str = "embedding",
     pdg_topn_map: "TopNMap | None" = None,
     mat_topn_map: "TopNMap | None" = None,
+    sec_type_topn_map: "TopNMap | None" = None,
     other_policy: str = "sample",
     seed: int | None = None,
     stage1_ddpm_steps: int = 1000,
@@ -469,14 +470,17 @@ def rollout(
     autoregressive) is inferred from `sec_decoder`'s own class — see
     `sample_stage1`/`sample_stage2` (giant.sample).
 
-    `pdg_topn_map`/`mat_topn_map` serve two independent purposes that happen
-    to share `pdg_topn_map` (one PDG map, not two): they're required
-    whenever `particle_conditioning`/`material_conditioning` is `"onehot"`
-    (feeds `build_cond_features`'s extra `cond_cat` top-N columns), and
-    `pdg_topn_map`/`other_policy` are additionally read under
+    `pdg_topn_map`/`mat_topn_map`/`sec_type_topn_map` serve three independent
+    purposes, no longer required to share one map (see gitea #29):
+    `pdg_topn_map`/`mat_topn_map` are required whenever
+    `particle_conditioning`/`material_conditioning` is `"onehot"` (feeds
+    `build_cond_features`'s extra `cond_cat` top-N columns); `sec_type_topn_map`/
+    `other_policy` are required instead under
     `stage2_model.particle_type.target = "onehot"` (secondary-species
-    decode). `seed` seeds the `other_policy = "sample"` draw only
-    (torch/numpy sampling itself is seeded by the caller, same as today).
+    decode) — its class count (`stage2_model.particle_type.n_classes`) may
+    differ from `pdg_topn_map`'s. `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 embedding-distance
     diagnostic across the whole run — see `L1DistCollector`. Only populated
@@ -487,6 +491,11 @@ def rollout(
             "conditioning.particle.type='onehot' rollout needs pdg_topn_map "
             "(the checkpoint's saved top-N map) — see ckpt['pdg_topn_map']"
         )
+    if sec_decoder.particle_type_cfg.get("target") == "onehot" and sec_type_topn_map is None:
+        raise RuntimeError(
+            "stage2_model.particle_type.target='onehot' rollout needs sec_type_topn_map "
+            "(the checkpoint's saved top-N map) — see ckpt['sec_type_topn_map']"
+        )
     if material_conditioning == "onehot" and mat_topn_map is None:
         raise RuntimeError(
             "conditioning.material.type='onehot' rollout needs mat_topn_map "
@@ -536,6 +545,7 @@ def rollout(
                     material_conditioning,
                     pdg_topn_map,
                     mat_topn_map,
+                    sec_type_topn_map,
                     other_policy,
                     rng,
                     stage1_ddpm_steps,
@@ -574,6 +584,7 @@ def _step_chunk(
     material_conditioning,
     pdg_topn_map,
     mat_topn_map,
+    sec_type_topn_map,
     other_policy,
     rng,
     stage1_ddpm_steps,
@@ -690,7 +701,7 @@ def _step_chunk(
         tr["pre_dir"],
         sec_phys_norm,
         pdg_map,
-        pdg_topn_map,
+        sec_type_topn_map,
         other_policy,
         rng,
     )

giant/training/loop.py

@@ -109,6 +109,7 @@ def train(
     mat_map: dict | None = None,
     proc_map: dict | None = None,
     pdg_topn_map: TopNMap | None = None,
+    sec_type_topn_map: TopNMap | None = None,
     mat_topn_map: TopNMap | None = None,
     model_config: dict | None = None,
     resume_path: str | Path | None = None,
@@ -142,6 +143,7 @@ def train(
         "mat_map": mat_map,
         "proc_map": proc_map,
         "pdg_topn_map": topnmap_to_json(pdg_topn_map) if pdg_topn_map is not None else None,
+        "sec_type_topn_map": topnmap_to_json(sec_type_topn_map) if sec_type_topn_map is not None else None,
         "mat_topn_map": topnmap_to_json(mat_topn_map) if mat_topn_map is not None else None,
         "model_config": model_config,
     }

giant/training/trainers.py

@@ -25,7 +25,7 @@ import torch.optim as optim
 from giant.config import ParticleTypeConfig, Stage1ModelConfig, Stage2ModelConfig, TrainConfig
 from giant.constants import CONT_SLOT_DIM
 from giant.data.dataset import StepBatch
-from giant.model.network import Router, stage2_type_dim
+from giant.model.network import Router, resolve_type_n_classes, stage2_type_dim
 from giant.model.schedule import (
     CosineSchedule,
     flow_matching_loss,
@@ -98,7 +98,7 @@ class StageSpec:
 
     # particle-type target (stage 2 only)
     particle_type: ParticleTypeConfig = field(default_factory=ParticleTypeConfig)
-    particle_type_emb_dim: int = 16
+    particle_type_n_classes: int = 16
 
     # optimization
     lr: float = 3e-4
@@ -151,7 +151,9 @@ class StageSpec:
             lambda_weight=stage_spec.lambda_weight,
             n_sec_lambda=s2_spec.n_sec.lambda_weight,
             particle_type=s2_spec.particle_type,
-            particle_type_emb_dim=cfg["conditioning"]["particle"]["emb_dim"],
+            particle_type_n_classes=resolve_type_n_classes(
+                s2_spec.particle_type.to_dict(), cfg["conditioning"]["particle"]["emb_dim"]
+            ),
             # train.* keys are all guaranteed by DEFAULT_CONFIG's deep-merge
             # (giant/config.py), so TrainConfig.from_dict never has to fall
             # back to a literal here; the field defaults below exist only
@@ -238,7 +240,7 @@ class StageTrainer:
         self._modules = (self.model, *extra_modules)
 
         self.particle_type_cfg = spec.particle_type.to_dict()
-        self.particle_type_emb_dim = spec.particle_type_emb_dim
+        self.particle_type_n_classes = spec.particle_type_n_classes
         self.ema_decay = spec.ema_decay
 
         self.ema_model: torch.nn.Module | None = None
@@ -329,7 +331,7 @@ class StageTrainer:
             n_sec,
             self.particle_type_cfg,
             self.model.cond_enc,
-            self.particle_type_emb_dim,
+            self.particle_type_n_classes,
             p_tf,
             self.spec.ar_sample_steps,
         )
@@ -355,7 +357,7 @@ class StageTrainer:
             self.particle_type_cfg,
             generator,
             self.model.cond_enc,
-            self.particle_type_emb_dim,
+            self.particle_type_n_classes,
         )
         return target.flatten(1) if flatten else target
 
@@ -756,7 +758,7 @@ class WGANStageTrainer(StageTrainer):
         multiplied on the fake side yet."""
         cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx = batch_tensors
         B = cond_cont.size(0)
-        type_dim = stage2_type_dim(self.particle_type_cfg, self.particle_type_emb_dim)
+        type_dim = stage2_type_dim(self.particle_type_cfg, self.particle_type_n_classes)
         slot_width = CONT_SLOT_DIM + type_dim
         k_max = sec_cont.size(1)
 
@@ -835,7 +837,7 @@ class WGANStageTrainer(StageTrainer):
                     fake_raw,
                     sec_cont.size(1),
                     CONT_SLOT_DIM,
-                    stage2_type_dim(self.particle_type_cfg, self.particle_type_emb_dim),
+                    stage2_type_dim(self.particle_type_cfg, self.particle_type_n_classes),
                     tau,
                     grad_probe=grad_probe,
                 )

tests/test_checkpoint_io.py

@@ -73,6 +73,16 @@ def _write_checkpoint(tmp_path, model_cfg=None, ema: bool = False, **ckpt_overri
     if ema:
         ckpt["model_ema"] = stage1.state_dict() if stage1 is not None else {}
         ckpt["sec_decoder_ema"] = stage2.state_dict() if stage2 is not None else {}
+    # DEFAULT_CONFIG's stage2_model.particle_type.target defaults to
+    # "onehot", and giant train's pipeline (gitea #29) now always writes a
+    # sec_type_topn_map in that case — default one in here too, unless a
+    # test explicitly overrides it, so fixtures represent a real, loadable
+    # checkpoint by default rather than exercising the "missing" guard by
+    # accident.
+    particle_type_target = cfg.get("stage2_model", {}).get("particle_type", {}).get("target", "onehot")
+    if particle_type_target == "onehot" and "sec_type_topn_map" not in ckpt_overrides:
+        default_sec_type_topn = TopNMap(class_map=dict(zip(PDG_MAP, range(len(PDG_MAP)))), other_members={})
+        ckpt["sec_type_topn_map"] = topnmap_to_json(default_sec_type_topn)
     ckpt.update(ckpt_overrides)
     path = tmp_path / "ckpt.pt"
     torch.save(ckpt, path)
@@ -179,6 +189,40 @@ def test_onehot_particle_conditioning_with_topn_map_succeeds(tmp_path):
     assert ctx.pdg_topn_map.class_map == {11: 0, 22: 1}
 
 
+def test_onehot_particle_type_target_without_sec_type_topn_map_raises(tmp_path):
+    """DEFAULT_CONFIG's stage2_model.particle_type.target="onehot" needs a
+    sec_type_topn_map (gitea #29) — a checkpoint with neither key at all
+    (not even the pre-#29 pdg_topn_map to fall back to) must fail loudly."""
+    checkpoint = _write_checkpoint(tmp_path, sec_type_topn_map=None)
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["sec_type_topn_map"]
+    torch.save(ckpt, checkpoint)
+
+    with pytest.raises(CheckpointCompatibilityError, match="sec_type_topn_map"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_pre_gitea_29_checkpoint_falls_back_to_pdg_topn_map_for_sec_type(tmp_path):
+    """A checkpoint written before gitea #29 has no sec_type_topn_map key at
+    all — conditioning and secondary-type onehot maps were always the same
+    map, saved once under pdg_topn_map. load_for_inference must reproduce
+    that exact pre-#29 behavior for such a checkpoint."""
+    topn = TopNMap(class_map={11: 0, 22: 1, -11: 2}, other_members={})
+    checkpoint = _write_checkpoint(
+        tmp_path,
+        model_cfg=_onehot_model_cfg(),
+        pdg_topn_map=topnmap_to_json(topn),
+        sec_type_topn_map=None,
+    )
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["sec_type_topn_map"]
+    torch.save(ckpt, checkpoint)
+
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+    assert ctx.sec_type_topn_map is not None
+    assert ctx.sec_type_topn_map.class_map == {11: 0, 22: 1, -11: 2}
+
+
 def test_ema_weights_requested_but_missing_raises(tmp_path):
     checkpoint = _write_checkpoint(tmp_path, ema=False)
 

tests/test_config.py

@@ -75,6 +75,16 @@ def test_stage2_model_config_defaults_match_documented_v030_intent():
     assert spec.particle_type.target == "onehot"
 
 
+def test_particle_type_config_n_classes_defaults_to_zero_and_round_trips():
+    """gitea #29: n_classes=0 means "inherit conditioning.particle.emb_dim"
+    — the default must stay 0 so an existing config.toml with no
+    stage2_model.particle_type.n_classes key reproduces pre-#29 behavior."""
+    assert gconfig.ParticleTypeConfig().n_classes == 0
+    spec = gconfig.ParticleTypeConfig.from_dict({"n_classes": 32})
+    assert spec.n_classes == 32
+    assert spec.to_dict()["n_classes"] == 32
+
+
 def test_router_config_extra_round_trips_composed_axis_keys():
     d = {"enabled": True, "type": "composed", "axis0_type": "energy", "axis0_n_experts": 4}
     router = gconfig.RouterConfig.from_dict(d)

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) ─
 
 

tests/test_pipeline.py

@@ -152,28 +152,60 @@ def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch):
     assert "normalizer: cache hit" in joined
 
 
-def test_run_train_job_builds_caches_and_persists_pdg_topn_map(tmp_path, data):
+def test_run_train_job_builds_caches_and_persists_sec_type_topn_map(tmp_path, data):
     """DEFAULT_CONFIG's stage2_model.particle_type.target defaults to
-    "onehot" — a plain _tiny_cfg() run must
-    build the shared pdg top-N map, cache it in the setup-cache sidecar, and
-    persist it into the checkpoint, with no extra config needed."""
+    "onehot" while conditioning.particle.type stays "physical" — a plain
+    _tiny_cfg() run must build the secondary-type-only pdg top-N map (gitea
+    #29: no longer shared with any conditioning-side onehot map), cache it in
+    the setup-cache sidecar, and persist it into the checkpoint's
+    sec_type_topn_map key, with no extra config needed. pdg_topn_map
+    (conditioning-only) stays unbuilt since conditioning.particle.type is
+    "physical" here."""
     echo1 = _run(data, tmp_path / "out1")
     assert any("building pdg top-N map" in m for m in echo1)
 
     loaded = setup_cache.load(data, [data])
     assert loaded is not None
-    key = setup_cache.topn_key("pdg", 4)  # conditioning.particle.emb_dim = 4
+    # stage2_model.particle_type.n_classes = 0 -> conditioning.particle.emb_dim = 4
+    key = setup_cache.topn_key("pdg", 4)
     assert key in loaded.topn_maps
     assert set(loaded.topn_maps[key].class_map.keys()) >= {11, 22}
 
     ckpt = torch.load(tmp_path / "out1" / "last.pt", weights_only=False)
-    assert "pdg_topn_map" in ckpt
-    assert set(ckpt["pdg_topn_map"]["class_map"].keys()) >= {"11", "22"}
+    assert ckpt.get("pdg_topn_map") is None
+    assert "sec_type_topn_map" in ckpt
+    assert set(ckpt["sec_type_topn_map"]["class_map"].keys()) >= {"11", "22"}
 
     echo2 = _run(data, tmp_path / "out2")
     assert any("pdg top-N map: cache hit" in m for m in echo2)
 
 
+def test_run_train_job_independent_cond_and_sec_type_topn_maps(tmp_path, data):
+    """conditioning.particle.type="onehot" and
+    stage2_model.particle_type.target="onehot" with different class counts
+    (gitea #29's fix: stage2_model.particle_type.n_classes decouples the two)
+    build two distinct top-N maps, cached under their own (axis, n_classes)
+    key and persisted under two distinct checkpoint keys — no longer forced
+    to share conditioning.particle.emb_dim."""
+    cfg = _tiny_cfg()
+    cfg["conditioning"]["particle"]["type"] = "onehot"  # emb_dim = 4, from _tiny_cfg
+    cfg["stage2_model"]["particle_type"]["n_classes"] = 3
+    echo = _run(data, tmp_path / "out", cfg=cfg)
+    assert any("mapped to 4 classes" in m for m in echo)
+    assert any("mapped to 3 classes" in m for m in echo)
+
+    loaded = setup_cache.load(data, [data])
+    assert loaded is not None
+    cond_key = setup_cache.topn_key("pdg", 4)
+    type_key = setup_cache.topn_key("pdg", 3)
+    assert cond_key in loaded.topn_maps
+    assert type_key in loaded.topn_maps
+
+    ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
+    assert ckpt.get("pdg_topn_map") is not None
+    assert ckpt.get("sec_type_topn_map") is not None
+
+
 def test_run_train_job_builds_caches_and_persists_material_topn_map(tmp_path, data):
     """conditioning.material.type="onehot" is an independent axis from the
     pdg one above, with its own build/cache-hit branch in run_setup_stage —

tests/test_rollout.py

@@ -430,7 +430,7 @@ def _run_v3(
     max_tracks_per_event=100,
     seeds=None,
     conditioning="physical",
-    pdg_topn_map=None,
+    sec_type_topn_map=None,
     other_policy="sample",
     seed=0,
     stage1_ddpm_steps=1000,
@@ -457,7 +457,7 @@ def _run_v3(
         escape_threshold=escape_threshold,
         particle_conditioning=conditioning,
         material_conditioning=conditioning,
-        pdg_topn_map=pdg_topn_map,
+        sec_type_topn_map=sec_type_topn_map,
         other_policy=other_policy,
         seed=seed,
         stage1_ddpm_steps=stage1_ddpm_steps,
@@ -532,7 +532,7 @@ def test_rollout_onehot_target_end_to_end(fake_material_props, decoder):
     (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")
+    rec = _run_v3(s1, s2, sec_type_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).
@@ -543,8 +543,8 @@ def test_rollout_onehot_target_end_to_end(fake_material_props, decoder):
 
 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)
+    with pytest.raises(RuntimeError, match="sec_type_topn_map"):
+        _run_v3(s1, s2, sec_type_topn_map=None)
 
 
 # --- conditioning.{particle,material}.type = "onehot" — a separate axis from
@@ -627,6 +627,86 @@ def test_rollout_conditioning_onehot_material_missing_topn_map_raises(
         _run_onehot_conditioning(mat_topn_map=None)
 
 
+SEC_TYPE_TOPN_MAP_DIFFERENT_N = TopNMap(class_map={22: 0, 11: 1, -11: 2, 13: 3}, other_members={2112: 3, 2212: 1})
+
+
+def _run_conditioning_and_type_onehot_different_n_classes():
+    """Both conditioning.particle.type="onehot" and
+    stage2_model.particle_type.target="onehot" active at once, with
+    stage2_model.particle_type.n_classes deliberately different from
+    conditioning.particle.emb_dim (gitea #29)."""
+    cond_emb_dim = len(PDG_MAP)  # 3
+    type_n_classes = 5  # deliberately different from cond_emb_dim
+    particle_cfg = {"type": "onehot", "emb_dim": cond_emb_dim, "n_layers": 1}
+    material_cfg = {"type": "onehot", "emb_dim": len(MAT_MAP), "n_layers": 1}
+    particle_type_cfg = {"target": "onehot", "n_classes": type_n_classes}
+
+    s1 = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+    ).eval()
+    sec_dim = stage2_trunk_sec_dim(particle_type_cfg, "flow", K_MAX, type_n_classes)
+    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=sec_dim,
+        generator="flow",
+        time_dim=16,
+        k_max=K_MAX,
+        particle_type_cfg=particle_type_cfg,
+    ).eval()
+    # Sanity: the model's own type_dim followed n_classes, not cond_emb_dim.
+    assert s2.type_dim == type_n_classes
+
+    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=15,
+        steps=3,
+        batch_size=128,
+        max_tracks_per_event=100,
+        escape_threshold=1e9,
+        particle_conditioning="onehot",
+        material_conditioning="onehot",
+        pdg_topn_map=COND_PDG_TOPN_MAP,
+        mat_topn_map=COND_MAT_TOPN_MAP,
+        sec_type_topn_map=SEC_TYPE_TOPN_MAP_DIFFERENT_N,
+        other_policy="modal",
+    )
+
+
+def test_rollout_conditioning_and_type_onehot_with_different_n_classes(fake_material_props):
+    """gitea #29 end-to-end: conditioning.particle.type="onehot" and
+    stage2_model.particle_type.target="onehot" now use independently sized
+    top-N maps (stage2_model.particle_type.n_classes != conditioning.particle
+    .emb_dim), and rollout must decode secondaries using the type-side map,
+    not silently reuse the conditioning-side one (the pre-#29 bug)."""
+    rec = _run_conditioning_and_type_onehot_different_n_classes()
+    assert len(rec["event_id"]) > 0
+    possible = set(SEC_TYPE_TOPN_MAP_DIFFERENT_N.class_map.keys()) | set(
+        SEC_TYPE_TOPN_MAP_DIFFERENT_N.other_members.keys()
+    )
+    secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())
+    assert secondary_pdgs <= possible
+
+
 @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