From 593c5f4d34385602f06e705f23e9da92ae99c067 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Fri, 14 Aug 2026 10:57:47 +0200 Subject: [PATCH] Deduplicate n_sec_head/type_head MLPs into build_mlp_head (gitea #36) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The same two-layer classifier head (Linear(cond_out_dim, hidden_dim // 2) -> SiLU -> Linear(hidden_dim // 2, out_dim)) was hand-rolled five times in giant/model/models.py: Stage1Model.n_sec_head, Stage2OneShot.n_sec_head/ .type_head, and Stage2Autoregressive.n_sec_head/.type_head. The `// 2` ratio and fixed 2-layer depth were undocumented magic numbers, and both n_sec accuracy and secondary-species accuracy are known weak spots that were untunable independently of the trunk they hang off. Adds `build_mlp_head(in_dim, out_dim, hidden, depth, act)` to giant/model/layers.py (depth=1 is a bare Linear; depth>=2 matches the old hardcoded shape exactly), and a new `HeadConfig` (hidden_ratio, depth) dataclass in giant/config.py, wired in as `stage1_model.heads.n_sec` and `stage2_model.heads.{n_sec,type}` — split per head type (not one shared block per stage) since n_sec and species prediction are called out as separate weak spots that may want independent capacity. Defaults (hidden_ratio=0.5, depth=2) reproduce the old hardcoded architecture bit-for-bit, so every existing config.toml and migrated v0.2 checkpoint is unaffected; no changes were needed to migrate_config or the legacy migration surfaces. No new CLI flags, matching how other nested sub-config (router.*, trunk.*) is set via config.toml rather than per-field flags. Co-Authored-By: Claude Opus 5 --- giant/config.py | 64 +++++++++++++++++++ giant/model/builders.py | 5 ++ giant/model/layers.py | 25 ++++++++ giant/model/models.py | 48 +++++++-------- giant/model/network.py | 2 + tests/test_config.py | 18 ++++++ tests/test_layers.py | 42 +++++++++++++ tests/test_network.py | 132 ++++++++++++++++++++++++++++++++++++++++ 8 files changed, 310 insertions(+), 26 deletions(-) create mode 100644 tests/test_layers.py diff --git a/giant/config.py b/giant/config.py index 1310c5e..b84d648 100644 --- a/giant/config.py +++ b/giant/config.py @@ -516,6 +516,64 @@ class AutoregressiveConfig: } +@dataclass(frozen=True) +class HeadConfig: + """A single classifier head's shape — `n_sec_head`/`type_head` (gitea + #36 deduplicated their five identical hand-rolled + `Linear -> SiLU -> Linear` definitions into + `giant.model.layers.build_mlp_head`, which this config drives). + `hidden_ratio=0.5`/`depth=2` are the exact pre-#36 hardcoded values + (hidden width = `hidden_dim // 2`, one hidden layer), so omitting a + `heads` block — including every migrated v0.2 config — reproduces the + old architecture bit-for-bit.""" + + hidden_ratio: float = 0.5 # hidden width = round(hidden_dim * hidden_ratio) + depth: int = 2 # matches build_mlp_head's depth + + @classmethod + def from_dict(cls, d: dict | None) -> "HeadConfig": + d = d or {} + return cls(hidden_ratio=d.get("hidden_ratio", 0.5), depth=d.get("depth", 2)) + + def to_dict(self) -> dict: + return {"hidden_ratio": self.hidden_ratio, "depth": self.depth} + + +@dataclass(frozen=True) +class Stage1HeadsConfig: + """Stage 1 only ever owns `n_sec_head`, and only for a migrated v0.2 + checkpoint (`stage2_model.n_sec.owner = "stage1"`) — see + `Stage1Model`'s docstring.""" + + n_sec: HeadConfig = field(default_factory=HeadConfig) + + @classmethod + def from_dict(cls, d: dict | None) -> "Stage1HeadsConfig": + d = d or {} + return cls(n_sec=HeadConfig.from_dict(d.get("n_sec"))) + + def to_dict(self) -> dict: + return {"n_sec": self.n_sec.to_dict()} + + +@dataclass(frozen=True) +class Stage2HeadsConfig: + """`n_sec` and `type` are independently configurable — n_sec accuracy + and secondary-species accuracy are separately known weak spots (gitea + #36).""" + + n_sec: HeadConfig = field(default_factory=HeadConfig) + type: HeadConfig = field(default_factory=HeadConfig) + + @classmethod + def from_dict(cls, d: dict | None) -> "Stage2HeadsConfig": + d = d or {} + return cls(n_sec=HeadConfig.from_dict(d.get("n_sec")), type=HeadConfig.from_dict(d.get("type"))) + + def to_dict(self) -> dict: + return {"n_sec": self.n_sec.to_dict(), "type": self.type.to_dict()} + + @dataclass(frozen=True) class Stage1ModelConfig: # false skips building/training stage 1 entirely. The resulting @@ -540,6 +598,7 @@ class Stage1ModelConfig: wgan: Stage1WganConfig = field(default_factory=Stage1WganConfig) router: RouterConfig = field(default_factory=RouterConfig) trunk: TrunkConfig = field(default_factory=TrunkConfig) + heads: Stage1HeadsConfig = field(default_factory=Stage1HeadsConfig) @classmethod def from_dict(cls, d: dict | None) -> "Stage1ModelConfig": @@ -556,6 +615,7 @@ class Stage1ModelConfig: wgan=Stage1WganConfig.from_dict(d.get("wgan")), router=RouterConfig.from_dict(d.get("router")), trunk=TrunkConfig.from_dict(d.get("trunk")), + heads=Stage1HeadsConfig.from_dict(d.get("heads")), ) def to_dict(self) -> dict: @@ -571,6 +631,7 @@ class Stage1ModelConfig: "wgan": self.wgan.to_dict(), "router": self.router.to_dict(), "trunk": self.trunk.to_dict(), + "heads": self.heads.to_dict(), } @@ -609,6 +670,7 @@ class Stage2ModelConfig: wgan: Stage2WganConfig = field(default_factory=Stage2WganConfig) router: Stage2RouterConfig = field(default_factory=Stage2RouterConfig) trunk: TrunkConfig = field(default_factory=TrunkConfig) + heads: Stage2HeadsConfig = field(default_factory=Stage2HeadsConfig) @classmethod def from_dict(cls, d: dict | None) -> "Stage2ModelConfig": @@ -632,6 +694,7 @@ class Stage2ModelConfig: wgan=Stage2WganConfig.from_dict(d.get("wgan")), router=Stage2RouterConfig.from_dict(d.get("router")), trunk=TrunkConfig.from_dict(d.get("trunk")), + heads=Stage2HeadsConfig.from_dict(d.get("heads")), ) def to_dict(self) -> dict: @@ -654,6 +717,7 @@ class Stage2ModelConfig: "wgan": self.wgan.to_dict(), "router": self.router.to_dict(), "trunk": self.trunk.to_dict(), + "heads": self.heads.to_dict(), } diff --git a/giant/model/builders.py b/giant/model/builders.py index 61c7318..21a3191 100644 --- a/giant/model/builders.py +++ b/giant/model/builders.py @@ -89,6 +89,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: block_conditioning=s1_spec.trunk.block_conditioning, n_sec_head_k_max=n_sec_head_k_max, cond_enc=shared_cond_enc, + n_sec_head_cfg=s1_spec.heads.n_sec.to_dict(), ) if s2_spec.active: @@ -134,6 +135,8 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: attn_n_heads=ar_cfg.attn_n_heads, attn_n_layers=ar_cfg.attn_n_layers, cond_enc=shared_cond_enc, + n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(), + type_head_cfg=s2_spec.heads.type.to_dict(), ) else: sec_dim = stage2_trunk_sec_dim( @@ -160,6 +163,8 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: build_n_sec_head=n_sec_owner != "stage1", particle_type_cfg=particle_type_cfg, cond_enc=shared_cond_enc, + n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(), + type_head_cfg=s2_spec.heads.type.to_dict(), ) return result diff --git a/giant/model/layers.py b/giant/model/layers.py index 64591ea..32c8aa1 100644 --- a/giant/model/layers.py +++ b/giant/model/layers.py @@ -42,6 +42,31 @@ def _make_axis_mlp(in_dim: int, emb_dim: int, n_layers: int) -> nn.Sequential: return nn.Sequential(*layers) +def build_mlp_head( + in_dim: int, out_dim: int, hidden: int, depth: int = 2, act: type[nn.Module] = nn.SiLU +) -> nn.Sequential: + """`depth`-layer MLP head (gitea #36) — factors out the n_sec_head/ + type_head pattern duplicated five times across `giant.model.models`. + + `depth=1` is a bare `Linear(in_dim, out_dim)` (no hidden layer/ + activation); `depth>=2` is `Linear(in_dim, hidden) -> act -> [Linear + (hidden, hidden) -> act] * (depth-2) -> Linear(hidden, out_dim)` — + `depth=2` reproduces every pre-#36 n_sec_head/type_head exactly when + `hidden == hidden_dim // 2`. Mirrors `_make_axis_mlp`'s depth + convention above, but takes `hidden` and `out_dim` as independent + widths (n_sec_head/type_head's hidden width is not their output width, + unlike the particle/material axis MLPs).""" + if depth < 1: + raise ValueError(f"depth must be >= 1, got {depth}") + if depth == 1: + return nn.Sequential(nn.Linear(in_dim, out_dim)) + layers: list[nn.Module] = [nn.Linear(in_dim, hidden), act()] + for _ in range(depth - 2): + layers += [nn.Linear(hidden, hidden), act()] + layers.append(nn.Linear(hidden, out_dim)) + return nn.Sequential(*layers) + + class ContextAdapter(nn.Module): """Projects a stage's outcome (e.g. Stage 1's 9D target) down to a fixed-width context vector for a downstream stage's conditioning — diff --git a/giant/model/models.py b/giant/model/models.py index 3896704..ed22a99 100644 --- a/giant/model/models.py +++ b/giant/model/models.py @@ -4,10 +4,11 @@ import torch import torch.nn as nn +from giant.config import HeadConfig from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, SEC_SLOT_DIM, X_DIM from giant.model.encoders import ConditionEncoder from giant.model.history import HistoryEncoder, build_history -from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding +from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding, build_mlp_head from giant.model.objectives import build_objective from giant.model.routers import Router from giant.model.trunks import build_trunk @@ -97,6 +98,7 @@ class Stage1Model(nn.Module): block_conditioning: str = "add", n_sec_head_k_max: int | None = None, cond_enc: ConditionEncoder | None = None, + n_sec_head_cfg: dict | None = None, ) -> None: super().__init__() self.generator_kind = generator @@ -116,11 +118,9 @@ class Stage1Model(nn.Module): ) self.n_sec_head = None if n_sec_head_k_max is not None: - self.n_sec_head = nn.Sequential( - nn.Linear(cond_out_dim, hidden_dim // 2), - nn.SiLU(), - nn.Linear(hidden_dim // 2, n_sec_head_k_max + 1), - ) + head_cfg = HeadConfig.from_dict(n_sec_head_cfg) + hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio)) + self.n_sec_head = build_mlp_head(cond_out_dim, n_sec_head_k_max + 1, hidden, head_cfg.depth) def forward( self, @@ -197,6 +197,8 @@ class Stage2OneShot(nn.Module): build_n_sec_head: bool = True, particle_type_cfg: dict | None = None, cond_enc: ConditionEncoder | None = None, + n_sec_head_cfg: dict | None = None, + type_head_cfg: dict | None = None, ) -> None: super().__init__() self.generator_kind = generator @@ -234,20 +236,16 @@ class Stage2OneShot(nn.Module): ) self.n_sec_head = None if build_n_sec_head: - self.n_sec_head = nn.Sequential( - nn.Linear(cond_out_dim, hidden_dim // 2), - nn.SiLU(), - nn.Linear(hidden_dim // 2, k_max + 1), - ) + head_cfg = HeadConfig.from_dict(n_sec_head_cfg) + hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio)) + self.n_sec_head = build_mlp_head(cond_out_dim, k_max + 1, hidden, head_cfg.depth) self.type_head = None target = self.particle_type_cfg.get("target", "physical") if target != "physical" and not objective.folds_type_slice: 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(), - nn.Linear(hidden_dim // 2, k_max * emb_dim), - ) + head_cfg = HeadConfig.from_dict(type_head_cfg) + hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio)) + self.type_head = build_mlp_head(cond_out_dim, k_max * emb_dim, hidden, head_cfg.depth) self._type_k_max = k_max self._type_emb_dim = emb_dim @@ -356,6 +354,8 @@ class Stage2Autoregressive(nn.Module): attn_n_heads: int = 4, attn_n_layers: int = 2, cond_enc: ConditionEncoder | None = None, + n_sec_head_cfg: dict | None = None, + type_head_cfg: dict | None = None, ) -> None: super().__init__() self.history_kind = history @@ -411,19 +411,15 @@ class Stage2Autoregressive(nn.Module): self.n_sec_head = None if build_n_sec_head: - self.n_sec_head = nn.Sequential( - nn.Linear(cond_out_dim, hidden_dim // 2), - nn.SiLU(), - nn.Linear(hidden_dim // 2, k_max + 1), - ) + head_cfg = HeadConfig.from_dict(n_sec_head_cfg) + hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio)) + self.n_sec_head = build_mlp_head(cond_out_dim, k_max + 1, hidden, head_cfg.depth) self.type_head = None target = self.particle_type_cfg.get("target", "physical") if target != "physical" and not objective.folds_type_slice: - self.type_head = nn.Sequential( - nn.Linear(cond_out_dim, hidden_dim // 2), - nn.SiLU(), - nn.Linear(hidden_dim // 2, self.type_dim), - ) + head_cfg = HeadConfig.from_dict(type_head_cfg) + hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio)) + self.type_head = build_mlp_head(cond_out_dim, self.type_dim, hidden, head_cfg.depth) def _base_cond(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor: base = self.cond_enc(cond_cont, cond_cat) diff --git a/giant/model/network.py b/giant/model/network.py index cb21523..65dd3b9 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -28,6 +28,7 @@ from giant.model.layers import ( SinusoidalEmbedding, _make_axis_mlp, build_block, + build_mlp_head, register_block, ) from giant.model.models import ( @@ -116,6 +117,7 @@ __all__ = [ "build_critics", "build_expert_body", "build_history", + "build_mlp_head", "build_models", "build_objective", "build_router", diff --git a/tests/test_config.py b/tests/test_config.py index ea82ba5..2945df6 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -55,6 +55,9 @@ def test_giant_config_to_dict_matches_default_config(): gconfig.NSecConfig, gconfig.ParticleTypeConfig, gconfig.AutoregressiveConfig, + gconfig.HeadConfig, + gconfig.Stage1HeadsConfig, + gconfig.Stage2HeadsConfig, gconfig.Stage1ModelConfig, gconfig.Stage2ModelConfig, gconfig.TrainConfig, @@ -94,6 +97,21 @@ def test_trunk_config_defaults_block_conditioning_to_add_for_both_stages(): assert gconfig.DEFAULT_CONFIG["stage2_model"]["trunk"]["block_conditioning"] == "add" +def test_heads_config_defaults_reproduce_pre_gitea_36_hardcoded_shape(): + """gitea #36: a pre-existing config with no `heads` key must reproduce + today's hardcoded `hidden_dim // 2`, one-hidden-layer architecture + exactly.""" + assert gconfig.Stage1ModelConfig().heads.n_sec.hidden_ratio == 0.5 + assert gconfig.Stage1ModelConfig().heads.n_sec.depth == 2 + assert gconfig.Stage2ModelConfig().heads.n_sec.hidden_ratio == 0.5 + assert gconfig.Stage2ModelConfig().heads.n_sec.depth == 2 + assert gconfig.Stage2ModelConfig().heads.type.hidden_ratio == 0.5 + assert gconfig.Stage2ModelConfig().heads.type.depth == 2 + assert gconfig.DEFAULT_CONFIG["stage1_model"]["heads"]["n_sec"] == {"hidden_ratio": 0.5, "depth": 2} + assert gconfig.DEFAULT_CONFIG["stage2_model"]["heads"]["n_sec"] == {"hidden_ratio": 0.5, "depth": 2} + assert gconfig.DEFAULT_CONFIG["stage2_model"]["heads"]["type"] == {"hidden_ratio": 0.5, "depth": 2} + + 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 diff --git a/tests/test_layers.py b/tests/test_layers.py new file mode 100644 index 0000000..a3f1e74 --- /dev/null +++ b/tests/test_layers.py @@ -0,0 +1,42 @@ +import pytest +import torch + +from giant.model.layers import build_mlp_head + + +def test_build_mlp_head_depth_1_is_bare_linear(): + head = build_mlp_head(8, 4, hidden=16, depth=1) + assert len(head) == 1 + assert isinstance(head[0], torch.nn.Linear) + assert head[0].in_features == 8 + assert head[0].out_features == 4 + out = head(torch.randn(3, 8)) + assert out.shape == (3, 4) + + +def test_build_mlp_head_depth_2_matches_pre_gitea_36_shape(): + head = build_mlp_head(8, 4, hidden=16, depth=2) + assert len(head) == 3 + assert isinstance(head[0], torch.nn.Linear) + assert head[0].in_features == 8 + assert head[0].out_features == 16 + assert isinstance(head[1], torch.nn.SiLU) + assert isinstance(head[2], torch.nn.Linear) + assert head[2].in_features == 16 + assert head[2].out_features == 4 + out = head(torch.randn(5, 8)) + assert out.shape == (5, 4) + + +def test_build_mlp_head_depth_3_has_extra_hidden_layer(): + head = build_mlp_head(8, 4, hidden=16, depth=3) + assert len(head) == 5 + widths = [(m.in_features, m.out_features) for m in head if isinstance(m, torch.nn.Linear)] + assert widths == [(8, 16), (16, 16), (16, 4)] + out = head(torch.randn(2, 8)) + assert out.shape == (2, 4) + + +def test_build_mlp_head_depth_0_raises(): + with pytest.raises(ValueError, match="depth"): + build_mlp_head(8, 4, hidden=16, depth=0) diff --git a/tests/test_network.py b/tests/test_network.py index f3fd5f2..d55aef9 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -93,6 +93,42 @@ def test_stage1_model_no_n_sec_head_by_default(): assert model.n_sec_head is None +def test_stage1_model_n_sec_head_default_cfg_matches_pre_gitea_36_shape(): + """No n_sec_head_cfg given must reproduce the old hardcoded + hidden_dim // 2, one-hidden-layer architecture exactly (gitea #36).""" + model = Stage1Model( + pdg_vocab=3, + mat_vocab=2, + particle_cfg=PARTICLE_CFG, + material_cfg=MATERIAL_CFG, + hidden_dim=40, + cond_out_dim=12, + n_sec_head_k_max=15, + ) + assert model.n_sec_head is not None + assert len(model.n_sec_head) == 3 + assert model.n_sec_head[0].in_features == 12 + assert model.n_sec_head[0].out_features == 20 # hidden_dim // 2 + assert model.n_sec_head[2].out_features == 16 # k_max + 1 + + +def test_stage1_model_n_sec_head_cfg_controls_hidden_width_and_depth(): + model = Stage1Model( + pdg_vocab=3, + mat_vocab=2, + particle_cfg=PARTICLE_CFG, + material_cfg=MATERIAL_CFG, + hidden_dim=32, + cond_out_dim=16, + n_sec_head_k_max=15, + n_sec_head_cfg={"hidden_ratio": 0.25, "depth": 1}, + ) + assert model.n_sec_head is not None + assert len(model.n_sec_head) == 1 + assert model.n_sec_head[0].in_features == 16 + assert model.n_sec_head[0].out_features == 16 + + # --- cat_col_layout / stage2_type_dim / stage2_trunk_sec_dim --------------- @@ -268,6 +304,38 @@ def test_stage2_oneshot_predict_type_raises_when_no_type_head(): pass +def test_stage2_oneshot_n_sec_head_and_type_head_cfg_control_hidden_width_and_depth(): + """gitea #36: n_sec_head_cfg/type_head_cfg are independently tunable.""" + k_max, emb_dim = 5, 6 + particle_cfg = {"type": "onehot", "emb_dim": emb_dim, "n_layers": 1} + sec_dim = stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim) + model = Stage2OneShot( + pdg_vocab=5, + mat_vocab=3, + particle_cfg=particle_cfg, + material_cfg=MATERIAL_CFG, + hidden_dim=40, + n_res_blocks=1, + cond_out_dim=12, + context_dim=8, + sec_dim=sec_dim, + generator="flow", + k_max=k_max, + particle_type_cfg={"target": "onehot", "lambda": 1.0}, + n_sec_head_cfg={"hidden_ratio": 0.25, "depth": 1}, + type_head_cfg={"hidden_ratio": 0.75, "depth": 2}, + ) + assert model.n_sec_head is not None + assert len(model.n_sec_head) == 1 + assert model.n_sec_head[0].in_features == 12 + assert model.n_sec_head[0].out_features == k_max + 1 + + assert model.type_head is not None + assert len(model.type_head) == 3 + assert model.type_head[0].out_features == 30 # round(40 * 0.75) + assert model.type_head[2].out_features == k_max * emb_dim + + def test_stage2_oneshot_forward_shape_onehot_wgan(): B, k_max, emb_dim = 4, 5, 6 model = _build_stage2("onehot", "wgan", emb_dim=emb_dim) @@ -518,6 +586,37 @@ def test_stage2_autoregressive_particle_type_n_classes_overrides_conditioning_em assert model.type_head[-1].out_features == 20 +def test_stage2_autoregressive_n_sec_head_and_type_head_cfg_control_hidden_width_and_depth(): + """gitea #36, 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} + model = Stage2Autoregressive( + pdg_vocab=5, + mat_vocab=3, + particle_cfg=particle_cfg, + material_cfg=MATERIAL_CFG, + hidden_dim=40, + n_res_blocks=1, + cond_out_dim=12, + context_dim=8, + generator="flow", + k_max=5, + particle_type_cfg=particle_type_cfg, + n_sec_head_cfg={"hidden_ratio": 0.25, "depth": 1}, + type_head_cfg={"hidden_ratio": 0.75, "depth": 2}, + ) + assert model.n_sec_head is not None + assert len(model.n_sec_head) == 1 + assert model.n_sec_head[0].in_features == 12 + assert model.n_sec_head[0].out_features == 6 # k_max + 1 + + assert model.type_head is not None + assert len(model.type_head) == 3 + assert model.type_head[0].out_features == 30 # round(40 * 0.75) + assert model.type_head[2].out_features == model.type_dim + + @pytest.mark.parametrize("target", ["physical", "onehot", "embedding"]) @pytest.mark.parametrize("generator", ["wgan", "flow"]) @pytest.mark.parametrize("history", ["markov", "attention"]) @@ -813,6 +912,39 @@ def test_build_models_omitted_decoder_and_particle_type_match_default_config(): assert built["stage2"].particle_type_cfg["target"] == "onehot" +def test_build_models_custom_heads_block_controls_head_shapes(): + """gitea #36: stage{1,2}_model.heads flows all the way from config dict + through build_models to the actual constructed head shapes.""" + cfg = _partial_model_config() + cfg["stage1_model"] = { + "active": True, + "hidden_dim": 40, + "n_res_blocks": 1, + "heads": {"n_sec": {"hidden_ratio": 0.25, "depth": 1}}, + } + cfg["stage2_model"]["decoder"] = "one_shot" + cfg["stage2_model"]["generator"] = "flow" # wgan folds the type slice; no separate type_head + cfg["stage2_model"]["n_sec"] = {"owner": "stage1"} + cfg["stage2_model"]["heads"] = { + "n_sec": {"hidden_ratio": 0.25, "depth": 1}, + "type": {"hidden_ratio": 0.75, "depth": 2}, + } + + built = build_models(cfg) + stage1, stage2 = built["stage1"], built["stage2"] + assert stage1 is not None + assert stage2 is not None + + assert stage1.n_sec_head is not None + assert len(stage1.n_sec_head) == 1 # owner=stage1, so stage1 builds it + assert stage2.n_sec_head is None # owner=stage1, so stage2 doesn't + + assert isinstance(stage2, Stage2OneShot) + assert stage2.type_head is not None + assert len(stage2.type_head) == 3 + assert stage2.type_head[0].out_features == 6 # round(8 * 0.75) + + def test_build_critics_omitted_particle_type_matches_default_config(): cfg = _partial_model_config() cfg["stage2_model"]["generator"] = "wgan"