Give Stage1Model/Stage2OneShot/Stage2Autoregressive a shared StageModel base (gitea #39) #56

Merged

lars merged 3 commits from fix/issue-39 into master 2026-08-14 15:16:03 +02:00

Conversation 0 Commits 3 Files Changed 22

+620 -311

3 Commits

Author	SHA1	Message	Date
lars	cc37a55183	Bump patch version to 0.3.1 ... Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-14 15:05:32 +02:00
lars	c4b12b5e7a	Pass ConditioningAxisConfig/ParticleTypeConfig themselves instead of raw dicts (gitea #38) ... build_models/build_critics parsed model_config into frozen dataclasses (ConditioningConfig, Stage2ModelConfig, ...) but then threw the parsed sub-objects away and passed the original raw dicts (conditioning["particle"], s2_spec.particle_type.to_dict()) down into ConditionEncoder/StageModel/etc, which re-read them with their own hardcoded .get(key, default) fallbacks — each an independent copy of a fact the dataclass already stated once. Worst instance: giant/training/trainers.py:236 converted an already-parsed ParticleTypeConfig back into a dict for no reason. Threads ConditioningAxisConfig (particle_cfg/material_cfg) and ParticleTypeConfig (particle_type_cfg) as the actual dataclass instances through every signature that used to type them dict: ConditionEncoder, StageModel/CriticModel, resolve_type_n_classes/stage2_type_dim/ stage2_trunk_sec_dim, giant/model/builders.py, giant/sample.py, giant/training/stage2_inputs.py, giant/training/trainers.py (StageSpec/ StageTrainer), giant/pipeline.py, giant/rollout.py, giant/validate.py — so ty now catches a misspelled field instead of it silently falling back. No config-schema change: config.toml/checkpoint model_config keep the same nested-dict shape; only what happens after the existing X.from_dict(...) parse changes. User-confirmed scope decision: both axes (particle_cfg/material_cfg and particle_type_cfg), not just the more heavily-duplicated particle_type_cfg axis, and not stopping at the two most literal parse-then-discard round trips — matching the issue's own proposal. Preserved-default decision: StageModel's particle_type_cfg=None sentinel (hit only by direct/test construction — build_models always passes an explicit particle_type) still resolves to ParticleTypeConfig(target= "physical"), not ParticleTypeConfig()'s own target="onehot" config-file default — switching it would have silently grown an unused, gradient-less type_head on every test that constructs Stage2OneShot/Stage2Autoregressive without particle_type_cfg=, breaking their "every param has a grad" checks. New tests in tests/test_network.py: ConditionEncoder/StageModel store the exact ConditioningAxisConfig/ParticleTypeConfig instance passed in (identity, not just equality) — no internal dict round-trip — and build_models's output carries real dataclass instances end to end, not the plain dicts it produced before this fix. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-14 15:03:55 +02:00
lars	1a3c907571	Give Stage1Model/Stage2OneShot/Stage2Autoregressive a shared StageModel base (gitea #39) ... Stage1Model, Stage2OneShot and Stage2Autoregressive each independently implemented ~90 near-identical lines of __init__ scaffolding: build-or-share cond_enc, particle_type_cfg normalisation, objective -> time_emb -> merged_cond_dim -> build_trunk, and the n_sec_head/type_head classifier heads (plus their identical RuntimeError guards). Now unblocked by #33 (trunk registry), #34 (block-conditioning registry) and #36 (build_mlp_head), which settled what belongs in the shared base. Adds StageModel(nn.Module) owning all of that: __init__ builds/shares cond_enc and normalises particle_type_cfg; _build_trunk_and_heads, called by each subclass after it sets up its own conditioning-assembly modules (cond_enc alone for Stage1Model, a context-fusion path for the two Stage2 classes), builds the objective/time embedding/trunk and the n_sec_head/type_head guarded by the shared _require_n_sec_head/ _require_type_head (Stage1Model overrides the n_sec guard since its message points at stage 2, not stage 1). Public __init__ signatures, attribute names, and forward/predict_* behaviour are unchanged. Verified with a pre/post state_dict-key-set diff against the pre-refactor classes (bit-identical) before writing this commit, plus new parametrized tests pinning each class's state_dict key set and the generator -> time_emb contract the base now owns. tests/test_migration_ v02_v03.py's existing bit-identical old-vs-new forward comparison and the rest of tests/test_network.py's per-class coverage pass unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-14 14:37:58 +02:00