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>
lars
32aa5a5f92