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>

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,

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",

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