Make config dataclasses the single source of truth for DEFAULT_CONFIG

DEFAULT_CONFIG and build_models/build_critics/StageSpec.from_config's
inline .get(key, default) fallbacks had already drifted: two keys
(stage2_model.decoder, stage2_model.particle_type.target) resolved
differently depending on whether a config dict came from
merge_cli_overrides (fully populated, correct) or was hand-built and
partial (fell back to stale v0.2-shaped literals). Introduce frozen
dataclasses (GiantConfig and its nested blocks) in giant/config.py as
the actual single declaration of every default; DEFAULT_CONFIG is now
generated from them instead of hand-maintained, and build_models,
build_critics, and StageSpec.from_config consume the dataclasses
instead of duplicating literal fallbacks, so this class of drift can't
recur. Router/n_sec sub-blocks keep an `extra` catch-all for their
genuinely dynamic keys (composed-router axes, runtime-seeded
centers_init, legacy_owner).

Fixing the fallback surfaced the same latent bug in two existing
partial-config callers that had been silently depending on it: a
test fixture in test_train.py and scripts/warm_setup_cache.py's
minimal cfg (now merged against DEFAULT_CONFIG instead of hand-rolled,
closing the gap for good). See issues.md Issue 1.

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

tests/test_config.py

@@ -27,6 +27,80 @@ def test_conditioning_enum_has_onehot():
     }
 
 
+# ---------------------------------------------------------------------------
+# Config dataclasses (issues.md Issue 1)
+# ---------------------------------------------------------------------------
+
+
+def test_giant_config_to_dict_matches_default_config():
+    """DEFAULT_CONFIG is generated from GiantConfig().to_dict() (not
+    hand-maintained), so the two cannot structurally drift apart — but this
+    pins the *equality* too, catching e.g. a stray in-place mutation of
+    DEFAULT_CONFIG added elsewhere after import."""
+    assert gconfig.GiantConfig().to_dict() == gconfig.DEFAULT_CONFIG
+
+
+@pytest.mark.parametrize(
+    "cls",
+    [
+        gconfig.ConditioningAxisConfig,
+        gconfig.ConditioningConfig,
+        gconfig.FlowConfig,
+        gconfig.DdpmConfig,
+        gconfig.Stage1WganConfig,
+        gconfig.Stage2WganConfig,
+        gconfig.RouterConfig,
+        gconfig.Stage2RouterConfig,
+        gconfig.NSecConfig,
+        gconfig.ParticleTypeConfig,
+        gconfig.AutoregressiveConfig,
+        gconfig.Stage1ModelConfig,
+        gconfig.Stage2ModelConfig,
+        gconfig.TrainConfig,
+        gconfig.GiantConfig,
+    ],
+)
+def test_config_dataclass_from_dict_round_trips_through_to_dict(cls):
+    assert cls.from_dict(cls().to_dict()) == cls()
+    assert cls.from_dict(None) == cls()
+
+
+def test_stage2_model_config_defaults_match_documented_v030_intent():
+    """The two keys issues.md Issue 1 found drifted between DEFAULT_CONFIG
+    and build_models/StageSpec.from_config's own .get(key, default)
+    fallbacks — pinned directly against the dataclass that is now their
+    shared single source of truth."""
+    spec = gconfig.Stage2ModelConfig()
+    assert spec.decoder == "autoregressive"
+    assert spec.particle_type.target == "onehot"
+
+
+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)
+    assert router.enabled is True
+    assert router.extra == {"axis0_type": "energy", "axis0_n_experts": 4}
+    assert router.to_dict()["axis0_type"] == "energy"
+
+
+def test_stage2_router_config_tie_to_stage1_not_leaked_into_extra():
+    router = gconfig.Stage2RouterConfig.from_dict({"tie_to_stage1": True})
+    assert router.tie_to_stage1 is True
+    assert "tie_to_stage1" not in router.extra
+
+
+def test_stage1_router_config_has_no_tie_to_stage1_key():
+    """Stage 1's router schema must not gain stage 2's tie_to_stage1 key —
+    that would change every future run's saved config.toml shape."""
+    assert "tie_to_stage1" not in gconfig.RouterConfig().to_dict()
+
+
+def test_n_sec_config_extra_round_trips_legacy_owner():
+    n_sec = gconfig.NSecConfig.from_dict({"mode": "head", "legacy_owner": "stage1"})
+    assert n_sec.legacy_owner == "stage1"
+    assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "legacy_owner": "stage1"}
+
+
 # ---------------------------------------------------------------------------
 # _deep_merge
 # ---------------------------------------------------------------------------

tests/test_network.py

@@ -12,6 +12,7 @@ from giant.model.network import (
     Stage1Model,
     Stage2Autoregressive,
     Stage2OneShot,
+    build_critics,
     build_models,
     cat_col_layout,
     stage2_trunk_sec_dim,
@@ -652,3 +653,58 @@ def test_build_models_share_stages_true_shared_params_are_in_both_stage_paramete
     assert shared_ids
     assert shared_ids <= {id(p) for p in stage1.parameters()}
     assert shared_ids <= {id(p) for p in stage2.parameters()}
+
+
+# ── build_models/build_critics: DEFAULT_CONFIG fallback drift (issues.md #1) ─
+
+
+def _partial_model_config() -> dict:
+    """A hand-built model_config that omits stage2_model.decoder and
+    stage2_model.particle_type — deliberately not derived from
+    DEFAULT_CONFIG, unlike _minimal_model_config above. Regression fixture
+    for issues.md Issue 1: build_models/build_critics/StageSpec.from_config's
+    own fallback defaults for these two keys must equal DEFAULT_CONFIG's
+    ("autoregressive" / "onehot"), not the old, now-wrong v0.2-shaped
+    ("one_shot" / "physical") literals that used to live in three separate
+    .get(key, default) call sites."""
+    return {
+        "pdg_vocab": 3,
+        "mat_vocab": 2,
+        "conditioning": {
+            "particle": {"type": "physical", "emb_dim": 4, "n_layers": 1},
+            "material": {"type": "physical", "emb_dim": 4, "n_layers": 1},
+        },
+        "stage1_model": {"active": False},
+        "stage2_model": {
+            "generator": "wgan",
+            "hidden_dim": 8,
+            "n_res_blocks": 1,
+            "k_max": 3,
+            # decoder and particle_type deliberately omitted
+        },
+    }
+
+
+def test_build_models_omitted_decoder_and_particle_type_match_default_config():
+    built = build_models(_partial_model_config())
+    assert isinstance(built["stage2"], Stage2Autoregressive)
+    assert built["stage2"].particle_type_cfg["target"] == "onehot"
+
+
+def test_build_critics_omitted_particle_type_matches_default_config():
+    cfg = _partial_model_config()
+    cfg["stage2_model"]["generator"] = "wgan"
+    onehot_critic = build_critics(cfg)["stage2"]
+    assert onehot_critic is not None
+    onehot_in_dim = onehot_critic.input_proj.in_features
+
+    cfg["stage2_model"]["particle_type"] = {"target": "physical"}
+    physical_critic = build_critics(cfg)["stage2"]
+    assert physical_critic is not None
+    physical_in_dim = physical_critic.input_proj.in_features
+
+    # onehot's per-slot type width is emb_dim classes vs. physical's fixed
+    # (log-mass, charge) pair — different unless emb_dim happens to be 2, so
+    # this also confirms the critic was actually built in onehot mode by
+    # default, not silently falling back to physical.
+    assert onehot_in_dim != physical_in_dim

tests/test_train.py

@@ -265,6 +265,12 @@ def _base_cfg():
             "k_max": K_MAX,
             "context_dim": 16,
             "n_sec": {"mode": "head", "lambda": 0.1},
+            # Explicit, not relying on the fallback default (which is
+            # "onehot", matching DEFAULT_CONFIG — see issues.md Issue 1):
+            # the "physical"-labelled cases below (and this fixture's own
+            # comment history) intend this as the base "physical" case,
+            # with "*_onehot"/"*_embedding" cases opting in explicitly.
+            "particle_type": {"target": "physical", "lambda": 1.0},
             "flow": {"time_dim": 16},
             "ddpm": {"time_dim": 16, "n_steps": 50},
             "wgan": {
@@ -566,6 +572,20 @@ def test_flow_stage_trainer_ddpm_not_implemented_for_stage2():
         FlowDDPMStageTrainer(spec, torch.nn.Linear(1, 1), torch.device("cpu"))
 
 
+def test_stage_spec_from_config_omitted_decoder_and_particle_type_match_default_config():
+    """Regression for issues.md Issue 1: StageSpec.from_config's own fallback
+    defaults for stage2_model.decoder/particle_type must equal
+    DEFAULT_CONFIG's ("autoregressive" / "onehot"), not the old, now-wrong
+    ("one_shot" / "physical") literals a .get(key, default) call used to
+    supply when a hand-built cfg omitted these keys."""
+    cfg = _base_cfg()
+    del cfg["stage2_model"]["decoder"]
+    del cfg["stage2_model"]["particle_type"]
+    spec = StageSpec.from_config(cfg, "stage2", is_stage2=True, steps_per_epoch=1)
+    assert spec.decoder == "autoregressive"
+    assert spec.particle_type.target == "onehot"
+
+
 # --- AR trainer wiring (v0.3.0 step 5) --------------------------------------
 
 
@@ -664,7 +684,7 @@ def test_wgan_onehot_one_shot_also_gets_grad_norm_instrumentation():
 
 
 def test_wgan_physical_omits_grad_norm_slice_columns():
-    cfg = _base_cfg()  # default stage2_model has no particle_type -> "physical"
+    cfg = _base_cfg()  # _base_cfg's stage2_model.particle_type.target is "physical"
     with tempfile.TemporaryDirectory() as tmp:
         out_dir = Path(tmp) / "run"
         _run_train(cfg, out_dir)