diff --git a/giant/model/models.py b/giant/model/models.py index ed22a99..d1f79c6 100644 --- a/giant/model/models.py +++ b/giant/model/models.py @@ -67,18 +67,136 @@ def stage2_trunk_sec_dim(particle_type_cfg: dict, generator: str, k_max: int, em return k_max * CONT_SLOT_DIM -class Stage1Model(nn.Module): - """Predicts the 9D primary post-step vector. No `n_sec_head` — fresh runs - move it to stage 2, except for a migrated v0.2 checkpoint - (`n_sec_head_k_max` given), where it stays attached here - since that's where its weights live and what conditioning it was trained - against (see `_migrate_legacy_model_config`). +class StageModel(nn.Module): + """Base owning the scaffolding common to `Stage1Model`, `Stage2OneShot`, + `Stage2Autoregressive` (gitea #39): build-or-share `cond_enc`, + `particle_type_cfg` normalisation, and — via `_build_trunk_and_heads`, + called by each subclass's `__init__` once its own conditioning-assembly + modules exist — the objective/time-embedding/trunk construction and the + `n_sec_head`/`type_head` classifier heads. A subclass supplies only its + own conditioning assembly (`Stage1Model` uses `cond_enc` directly; + `Stage2OneShot`/`Stage2Autoregressive` add a context-fusion path) and its + trunk's output width. `cond_enc`, if given, is used in place of building a fresh `ConditionEncoder` — `conditioning.share_stages = true`: `build_models` constructs one shared instance and passes it to both stages, halving the conditioning parameter count and forcing a common representation.""" + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + particle_cfg: dict, + material_cfg: dict, + cond_out_dim: int, + generator: str, + noise_dim: int, + k_max: int | None = None, + particle_type_cfg: dict | None = None, + cond_enc: ConditionEncoder | None = None, + ) -> None: + super().__init__() + self.generator_kind = generator + 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, resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"]) + ) + self.cond_enc = ( + cond_enc + if cond_enc is not None + else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) + ) + + def _build_trunk_and_heads( + self, + *, + trunk_out_dim: int, + hidden_dim: int, + n_res_blocks: int, + cond_out_dim: int, + time_dim: int, + router: Router | None, + trunk_type: str, + block_conditioning: str, + dropout: float, + n_sec_head_k_max: int | None, + n_sec_head_cfg: dict | None, + type_head_out_dim: int | None, + type_head_cfg: dict | None, + ) -> None: + """Builds `self.time_emb`, `self.trunk`, `self.n_sec_head`, + `self.type_head`. Called by a subclass's `__init__` after it has set + up its own conditioning-assembly modules — `merged_cond_dim` below + must match the width that assembly (`_cond_embed`/`_base_cond`/ + `_token_cond`, or plain `cond_enc` for `Stage1Model`) actually + produces. + + `n_sec_head` is built iff `n_sec_head_k_max is not None` (output + width `n_sec_head_k_max + 1`) — `Stage1Model` passes this only for a + migrated v0.2 checkpoint, `Stage2OneShot`/`Stage2Autoregressive` pass + it whenever `build_n_sec_head=True`. `type_head` is built iff + `type_head_out_dim is not None` (the caller — only the two Stage2 + classes — passes `None` exactly when `particle_type_cfg["target"] == + "physical"`) *and* the objective doesn't fold the type slice into its + own trunk output (checked here, since `objective` is already needed + for the trunk itself). + """ + objective = build_objective(self.generator_kind) + has_time = objective.needs_time + self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None + merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim + in_dim = objective.trunk_in_dim(trunk_out_dim, self.noise_dim) + self.trunk = build_trunk( + router, + trunk_type, + in_dim, + trunk_out_dim, + hidden_dim, + n_res_blocks, + merged_cond_dim, + dropout, + block_conditioning, + ) + self.n_sec_head = None + if n_sec_head_k_max is not None: + 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) + self.type_head = None + if type_head_out_dim is not None and not objective.folds_type_slice: + 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, type_head_out_dim, hidden, head_cfg.depth) + + def _require_n_sec_head(self) -> None: + if self.n_sec_head is None: + raise RuntimeError( + f"this {type(self).__name__} has no n_sec_head — it belongs to " + "a migrated v0.2 checkpoint (n_sec.owner='stage1'); call " + "stage1.predict_n_sec(cond_cont, cond_cat) instead" + ) + + def _require_type_head(self) -> None: + if self.type_head is None: + raise RuntimeError( + f"this {type(self).__name__} has no type_head — either " + "particle_type.target='physical' (the type slice is part of " + "forward()'s own output) or generator='wgan' (the WGAN " + "trainer reads the type slice out of forward()'s output " + "directly instead)" + ) + + +class Stage1Model(StageModel): + """Predicts the 9D primary post-step vector. No `n_sec_head` — fresh runs + move it to stage 2, except for a migrated v0.2 checkpoint + (`n_sec_head_k_max` given), where it stays attached here + since that's where its weights live and what conditioning it was trained + against (see `_migrate_legacy_model_config`).""" + def __init__( self, pdg_vocab: int, @@ -100,27 +218,31 @@ class Stage1Model(nn.Module): cond_enc: ConditionEncoder | None = None, n_sec_head_cfg: dict | None = None, ) -> None: - super().__init__() - self.generator_kind = generator - self.noise_dim = noise_dim - self.cond_enc = ( - cond_enc - if cond_enc is not None - else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) + super().__init__( + pdg_vocab, + mat_vocab, + particle_cfg, + material_cfg, + cond_out_dim=cond_out_dim, + generator=generator, + noise_dim=noise_dim, + cond_enc=cond_enc, ) - objective = build_objective(generator) - has_time = objective.needs_time - self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None - merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim - in_dim = objective.trunk_in_dim(x_dim, noise_dim) - self.trunk = build_trunk( - router, trunk_type, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout, block_conditioning + self._build_trunk_and_heads( + trunk_out_dim=x_dim, + hidden_dim=hidden_dim, + n_res_blocks=n_res_blocks, + cond_out_dim=cond_out_dim, + time_dim=time_dim, + router=router, + trunk_type=trunk_type, + block_conditioning=block_conditioning, + dropout=dropout, + n_sec_head_k_max=n_sec_head_k_max, + n_sec_head_cfg=n_sec_head_cfg, + type_head_out_dim=None, + type_head_cfg=None, ) - self.n_sec_head = None - if n_sec_head_k_max is not None: - 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, @@ -133,21 +255,28 @@ class Stage1Model(nn.Module): cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb return self.trunk(x_t, cond, cond_cont, cond_cat) - def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - """Return n_sec logits (B, K_MAX+1) from conditioning alone. Only - valid on a migrated v0.2 checkpoint's Stage1Model — fresh v0.3.0 - configs predict n_sec from Stage2OneShot instead.""" + def _require_n_sec_head(self) -> None: + """Overrides `StageModel`'s guard — a `Stage1Model` with no + `n_sec_head` points the caller to stage 2 (n_sec's default owner), + not to `stage1` as the base's message would.""" if self.n_sec_head is None: raise RuntimeError( "this Stage1Model has no n_sec_head — n_sec now lives on " "stage 2 by default; this method only exists " "for a migrated v0.2 checkpoint (n_sec.owner='stage1')" ) + + def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + """Return n_sec logits (B, K_MAX+1) from conditioning alone. Only + valid on a migrated v0.2 checkpoint's Stage1Model — fresh v0.3.0 + configs predict n_sec from Stage2OneShot instead.""" + self._require_n_sec_head() + assert self.n_sec_head is not None c_emb = self.cond_enc(cond_cont, cond_cat) return self.n_sec_head(c_emb) -class Stage2OneShot(nn.Module): +class Stage2OneShot(StageModel): """Predicts all `k_max` secondary slots simultaneously — v0.2 behaviour, reproduced exactly (`decoder = "autoregressive"` is `Stage2Autoregressive`, step 4/5, not implemented yet). @@ -169,9 +298,6 @@ class Stage2OneShot(nn.Module): Under a folding objective (wgan) the type slice stays folded into `sec_dim` (just `emb_dim` instead of `PARTICLE_PHYS_DIM` wide) and `type_head` is unused (`None`) — the WGAN trainer handles the ST-Gumbel relaxation. - - `cond_enc`, if given, is used in place of building a fresh - `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`). """ def __init__( @@ -200,54 +326,40 @@ class Stage2OneShot(nn.Module): n_sec_head_cfg: dict | None = None, type_head_cfg: dict | None = None, ) -> None: - super().__init__() - self.generator_kind = generator - 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, resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"]) - ) - self.cond_enc = ( - cond_enc - if cond_enc is not None - else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) + super().__init__( + pdg_vocab, + mat_vocab, + particle_cfg, + material_cfg, + cond_out_dim=cond_out_dim, + generator=generator, + noise_dim=noise_dim, + k_max=k_max, + particle_type_cfg=particle_type_cfg, + cond_enc=cond_enc, ) self.context_adapter = ContextAdapter(x_dim, context_dim) self.fuse = nn.Sequential( nn.Linear(cond_out_dim + context_dim, cond_out_dim), nn.SiLU(), ) - objective = build_objective(generator) - has_time = objective.needs_time - self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None - merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim - in_dim = objective.trunk_in_dim(sec_dim, noise_dim) - self.trunk = build_trunk( - router, - trunk_type, - in_dim, - sec_dim, - hidden_dim, - n_res_blocks, - merged_cond_dim, - dropout, - block_conditioning, - ) - self.n_sec_head = None - if build_n_sec_head: - 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"]) - 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 + type_head_out_dim = None if target == "physical" else k_max * self.type_dim + self._build_trunk_and_heads( + trunk_out_dim=sec_dim, + hidden_dim=hidden_dim, + n_res_blocks=n_res_blocks, + cond_out_dim=cond_out_dim, + time_dim=time_dim, + router=router, + trunk_type=trunk_type, + block_conditioning=block_conditioning, + dropout=dropout, + n_sec_head_k_max=k_max if build_n_sec_head else None, + n_sec_head_cfg=n_sec_head_cfg, + type_head_out_dim=type_head_out_dim, + type_head_cfg=type_head_cfg, + ) def _cond_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor: base = self.cond_enc(cond_cont, cond_cat) @@ -272,12 +384,8 @@ class Stage2OneShot(nn.Module): cond_cat: torch.Tensor, stage1_out: torch.Tensor, ) -> torch.Tensor: - if self.n_sec_head is None: - raise RuntimeError( - "this Stage2OneShot has no n_sec_head — it belongs to a " - "migrated v0.2 checkpoint (n_sec.owner='stage1'); call " - "stage1.predict_n_sec(cond_cont, cond_cat) instead" - ) + self._require_n_sec_head() + assert self.n_sec_head is not None c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out) return self.n_sec_head(c_emb) @@ -291,19 +399,13 @@ class Stage2OneShot(nn.Module): vectors (`target="embedding"`) — only under `generator in ("flow", "ddpm")`; `generator == "wgan"` folds the type slice into `forward`'s own output instead (see class docstring).""" - if self.type_head is None: - raise RuntimeError( - "this Stage2OneShot has no type_head — either " - "particle_type.target='physical' (the type slice is part of " - "forward()'s own output) or generator='wgan' (the WGAN " - "trainer reads the type slice out of forward()'s output " - "directly instead)" - ) + self._require_type_head() + assert self.type_head is not None c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out) - return self.type_head(c_emb).view(-1, self._type_k_max, self._type_emb_dim) + return self.type_head(c_emb).view(-1, self.k_max, self.type_dim) -class Stage2Autoregressive(nn.Module): +class Stage2Autoregressive(StageModel): """Emits secondaries one at a time in descending-energy order, instead of `Stage2OneShot`'s simultaneous k_max-slot prediction. `history` selects `MarkovHistory` or @@ -324,9 +426,6 @@ class Stage2Autoregressive(nn.Module): on token position; `_token_cond` additionally fuses in the history encoding and two running scalars (remaining energy-budget fraction, normalized slot index), and feeds `forward`/`predict_type`/the trunk. - - `cond_enc`, if given, is used in place of building a fresh - `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`). """ def __init__( @@ -357,20 +456,19 @@ class Stage2Autoregressive(nn.Module): n_sec_head_cfg: dict | None = None, type_head_cfg: dict | None = None, ) -> None: - super().__init__() - self.history_kind = history - self.generator_kind = generator - self.noise_dim = noise_dim - self.k_max = k_max - self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"}) - 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 = ( - cond_enc - if cond_enc is not None - else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) + super().__init__( + pdg_vocab, + mat_vocab, + particle_cfg, + material_cfg, + cond_out_dim=cond_out_dim, + generator=generator, + noise_dim=noise_dim, + k_max=k_max, + particle_type_cfg=particle_type_cfg, + cond_enc=cond_enc, ) + self.history_kind = history self.context_adapter = ContextAdapter(x_dim, context_dim) self.base_fuse = nn.Sequential( nn.Linear(cond_out_dim + context_dim, cond_out_dim), @@ -391,35 +489,28 @@ class Stage2Autoregressive(nn.Module): nn.SiLU(), ) - objective = build_objective(generator) - has_time = objective.needs_time - self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None - merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim - token_dim = stage2_trunk_sec_dim(self.particle_type_cfg, generator, 1, emb_dim) - in_dim = objective.trunk_in_dim(token_dim, noise_dim) - self.trunk = build_trunk( - router, - trunk_type, - in_dim, - token_dim, - hidden_dim, - n_res_blocks, - merged_cond_dim, - dropout, - block_conditioning, - ) - - self.n_sec_head = None - if build_n_sec_head: - 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 + # `self.type_dim` (set by StageModel.__init__) doubles as the raw + # `emb_dim` `stage2_trunk_sec_dim` wants: for a non-"physical" target + # `stage2_type_dim` already resolved `type_dim` to exactly that value; + # for "physical" the emb_dim argument goes unused anyway. + token_dim = stage2_trunk_sec_dim(self.particle_type_cfg, generator, 1, self.type_dim) target = self.particle_type_cfg.get("target", "physical") - if target != "physical" and not objective.folds_type_slice: - 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) + type_head_out_dim = None if target == "physical" else self.type_dim + self._build_trunk_and_heads( + trunk_out_dim=token_dim, + hidden_dim=hidden_dim, + n_res_blocks=n_res_blocks, + cond_out_dim=cond_out_dim, + time_dim=time_dim, + router=router, + trunk_type=trunk_type, + block_conditioning=block_conditioning, + dropout=dropout, + n_sec_head_k_max=k_max if build_n_sec_head else None, + n_sec_head_cfg=n_sec_head_cfg, + type_head_out_dim=type_head_out_dim, + type_head_cfg=type_head_cfg, + ) 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) @@ -511,12 +602,8 @@ class Stage2Autoregressive(nn.Module): return out.view(B, K, -1) def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor: - if self.n_sec_head is None: - raise RuntimeError( - "this Stage2Autoregressive has no n_sec_head — it belongs to " - "a migrated v0.2 checkpoint (n_sec.owner='stage1'); call " - "stage1.predict_n_sec(cond_cont, cond_cat) instead" - ) + self._require_n_sec_head() + assert self.n_sec_head is not None return self.n_sec_head(self._base_cond(cond_cont, cond_cat, stage1_out)) def predict_type( @@ -530,14 +617,8 @@ class Stage2Autoregressive(nn.Module): slot_idx: torch.Tensor, hist: torch.Tensor | None = None, ) -> torch.Tensor: - if self.type_head is None: - raise RuntimeError( - "this Stage2Autoregressive has no type_head — either " - "particle_type.target='physical' (the type slice is part of " - "forward()'s own output) or generator='wgan' (the WGAN " - "trainer reads the type slice out of forward()'s output " - "directly instead)" - ) + self._require_type_head() + assert self.type_head is not None c_emb = self._token_cond( cond_cont, cond_cat, diff --git a/giant/model/network.py b/giant/model/network.py index 547b39b..ffd037a 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -36,6 +36,7 @@ from giant.model.models import ( Stage1Model, Stage2Autoregressive, Stage2OneShot, + StageModel, resolve_type_n_classes, stage2_trunk_sec_dim, stage2_type_dim, @@ -102,6 +103,7 @@ __all__ = [ "Stage1Model", "Stage2Autoregressive", "Stage2OneShot", + "StageModel", "TRUNK_REGISTRY", "Trunk", "WganObjective", diff --git a/tests/test_network.py b/tests/test_network.py index 4a99ad7..73662eb 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -14,9 +14,11 @@ from giant.model.network import ( Stage1Model, Stage2Autoregressive, Stage2OneShot, + StageModel, build_critics, build_history, build_models, + build_objective, stage2_trunk_sec_dim, stage2_type_dim, ) @@ -987,3 +989,163 @@ def test_build_critics_stage2_critic_hidden_dim_and_n_res_blocks_override_genera assert overridden is not None assert overridden.input_proj.out_features == 16 assert len(overridden.blocks) == 3 + + +# ── StageModel base (gitea #39): Stage1Model/Stage2OneShot/Stage2Autoregressive +# scaffolding — construction order, and therefore fresh-init RNG draw order and +# state_dict key set, must stay byte-for-byte what it was before the base class +# existed. ------------------------------------------------------------------ + +_STAGE_HIDDEN_DIM = 32 +_STAGE_N_BLOCKS = 2 +_STAGE_COND_OUT_DIM = 16 + + +def _resblock_keys(prefix: str) -> set[str]: + return { + f"{prefix}.norm.weight", + f"{prefix}.norm.bias", + f"{prefix}.linear1.weight", + f"{prefix}.linear1.bias", + f"{prefix}.cond_proj.weight", + f"{prefix}.linear2.weight", + f"{prefix}.linear2.bias", + } + + +def _trunk_keys(prefix: str = "trunk") -> set[str]: + keys = { + f"{prefix}.input_proj.weight", + f"{prefix}.input_proj.bias", + f"{prefix}.out_proj.weight", + f"{prefix}.out_proj.bias", + } + for i in range(_STAGE_N_BLOCKS): + keys |= _resblock_keys(f"{prefix}.blocks.{i}") + return keys + + +def _cond_enc_keys() -> set[str]: + return { + "cond_enc.mlp.0.weight", + "cond_enc.mlp.0.bias", + "cond_enc.mlp.2.weight", + "cond_enc.mlp.2.bias", + "cond_enc.particle_mlp.0.weight", + "cond_enc.particle_mlp.0.bias", + "cond_enc.material_mlp.0.weight", + "cond_enc.material_mlp.0.bias", + } + + +def _fuse_keys(name: str) -> set[str]: + return {f"{name}.0.weight", f"{name}.0.bias"} + + +def _head_keys(name: str) -> set[str]: + return {f"{name}.0.weight", f"{name}.0.bias", f"{name}.2.weight", f"{name}.2.bias"} + + +def _expected_stage_keys(*, has_time: bool, extra: set[str]) -> set[str]: + keys = _cond_enc_keys() | _trunk_keys() | extra + if has_time: + keys.add("time_emb.freqs") + return keys + + +@pytest.mark.parametrize("generator", ["flow", "wgan"]) +def test_stage1_model_state_dict_keys_unchanged_by_stagemodel_refactor(generator): + model = Stage1Model( + pdg_vocab=5, + mat_vocab=3, + particle_cfg=PARTICLE_CFG, + material_cfg=MATERIAL_CFG, + hidden_dim=_STAGE_HIDDEN_DIM, + n_res_blocks=_STAGE_N_BLOCKS, + cond_out_dim=_STAGE_COND_OUT_DIM, + generator=generator, + time_dim=8, + noise_dim=8, + n_sec_head_k_max=15, + ) + expected = _expected_stage_keys( + has_time=build_objective(generator).needs_time, + extra=_head_keys("n_sec_head"), + ) + assert set(model.state_dict().keys()) == expected + + +@pytest.mark.parametrize("generator", ["flow", "wgan"]) +def test_stage2_oneshot_state_dict_keys_unchanged_by_stagemodel_refactor(generator): + model = Stage2OneShot( + pdg_vocab=5, + mat_vocab=3, + particle_cfg=PARTICLE_CFG, + material_cfg=MATERIAL_CFG, + hidden_dim=_STAGE_HIDDEN_DIM, + n_res_blocks=_STAGE_N_BLOCKS, + cond_out_dim=_STAGE_COND_OUT_DIM, + generator=generator, + time_dim=8, + noise_dim=8, + k_max=15, + ) + extra = _head_keys("n_sec_head") | {"context_adapter.proj.weight", "context_adapter.proj.bias"} | _fuse_keys("fuse") + expected = _expected_stage_keys(has_time=build_objective(generator).needs_time, extra=extra) + assert set(model.state_dict().keys()) == expected + + +@pytest.mark.parametrize("generator", ["flow", "wgan"]) +def test_stage2_autoregressive_state_dict_keys_unchanged_by_stagemodel_refactor(generator): + model = Stage2Autoregressive( + pdg_vocab=5, + mat_vocab=3, + particle_cfg=PARTICLE_CFG, + material_cfg=MATERIAL_CFG, + hidden_dim=_STAGE_HIDDEN_DIM, + n_res_blocks=_STAGE_N_BLOCKS, + cond_out_dim=_STAGE_COND_OUT_DIM, + generator=generator, + time_dim=8, + noise_dim=8, + k_max=15, + ) + extra = ( + _head_keys("n_sec_head") + | {"context_adapter.proj.weight", "context_adapter.proj.bias"} + | _fuse_keys("base_fuse") + | _fuse_keys("token_fuse") + | {"history_encoder.start", "history_encoder.mlp.0.weight", "history_encoder.mlp.0.bias"} + ) + expected = _expected_stage_keys(has_time=build_objective(generator).needs_time, extra=extra) + assert set(model.state_dict().keys()) == expected + + +@pytest.mark.parametrize("cls", [Stage1Model, Stage2OneShot, Stage2Autoregressive]) +def test_stage_classes_are_stagemodel_subclasses(cls): + assert issubclass(cls, StageModel) + + +@pytest.mark.parametrize("cls", [Stage1Model, Stage2OneShot, Stage2Autoregressive]) +@pytest.mark.parametrize("generator", ["flow", "ddpm", "wgan"]) +def test_stagemodel_time_emb_matches_objective_needs_time(cls, generator): + kwargs = dict( + pdg_vocab=5, + mat_vocab=3, + particle_cfg=PARTICLE_CFG, + material_cfg=MATERIAL_CFG, + hidden_dim=_STAGE_HIDDEN_DIM, + n_res_blocks=_STAGE_N_BLOCKS, + cond_out_dim=_STAGE_COND_OUT_DIM, + generator=generator, + time_dim=8, + noise_dim=8, + ) + if cls is Stage1Model: + kwargs["n_sec_head_k_max"] = 15 + else: + kwargs["k_max"] = 15 + model = cls(**kwargs) + assert model.generator_kind == generator + assert model.noise_dim == 8 + assert (model.time_emb is not None) == build_objective(generator).needs_time