Make ResBlock's conditioning-injection mechanism selectable (gitea #34)

ResBlock injected conditioning exactly one way — h = linear1(h) +
cond_proj(cond), a conditional bias, the weakest standard option for a
model whose entire job is to be conditional. Adds BLOCK_REGISTRY
(giant/model/layers.py), mirroring the TRUNK_REGISTRY/ROUTER_REGISTRY
registry+factory idiom (gitea #33), with two new drop-in alternatives:
FilmResBlock (per-channel scale+shift modulating the norm output,
zero-init so conditioning has no effect at construction) and
AdaLNResBlock (DiT-style AdaLN-Zero — the norm's own affine is replaced
by a conditioning-derived scale/shift, plus a zero-init gate on the
residual branch, making the block the exact identity function at init).

Selected per stage via a new stage{1,2}_model.trunk.block_conditioning
config leaf ("add" | "film" | "adaln", default "add"), threaded through
build_trunk/build_expert_body/RoutedTrunk and the three stage model
constructors. Default stays "add" and ResBlock's body is unchanged, so
existing configs/checkpoints are bit-identical to before this change.

Decided during planning: the new field lives on the existing TrunkConfig
rather than a new top-level block/blocks config section; the WGAN
CriticModel (which builds its own ResBlock stack outside TRUNK_REGISTRY)
and the issue's mentioned blocks.norm/blocks.activation axes are both
left out of scope.

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

giant/config.py

@@ -346,17 +346,23 @@ class TrunkConfig:
     stage, unaffected by this block. A future body's own hyperparameters
     (e.g. a transformer's `n_heads`/`n_layers`) would get their own sibling
     field here, matching how `flow`/`ddpm`/`wgan` already coexist selected by
-    `generator`."""
+    `generator`.
+
+    `block_conditioning` selects each body's conditioning-injection mechanism
+    from `giant.model.layers.BLOCK_REGISTRY` — `"add"` (default, today's
+    conditional-bias `ResBlock`, bit-identical to pre-gitea-#34 behaviour),
+    `"film"`, or `"adaln"`."""
 
     type: str = "resmlp"
+    block_conditioning: str = "add"
 
     @classmethod
     def from_dict(cls, d: dict | None) -> "TrunkConfig":
         d = d or {}
-        return cls(type=d.get("type", "resmlp"))
+        return cls(type=d.get("type", "resmlp"), block_conditioning=d.get("block_conditioning", "add"))
 
     def to_dict(self) -> dict:
-        return {"type": self.type}
+        return {"type": self.type, "block_conditioning": self.block_conditioning}
 
 
 @dataclass(frozen=True)
@@ -1431,6 +1437,8 @@ _OUT_DIR_NAME_CANDIDATES = [
     ),
     ("stage1_trunk_type", _path_candidate("stage1_model.trunk.type", "s1t-")),
     ("stage2_trunk_type", _path_candidate("stage2_model.trunk.type", "s2t-")),
+    ("stage1_block_cond", _path_candidate("stage1_model.trunk.block_conditioning", "s1bc-")),
+    ("stage2_block_cond", _path_candidate("stage2_model.trunk.block_conditioning", "s2bc-")),
     ("stage1_router", _router_candidate("stage1_model", "s1")),
     ("stage2_router", _router_candidate("stage2_model", "s2")),
     (

giant/model/builders.py

@@ -84,6 +84,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
             noise_dim=s1_spec.wgan.noise_dim,
             router=stage1_router,
             trunk_type=s1_spec.trunk.type,
+            block_conditioning=s1_spec.trunk.block_conditioning,
             n_sec_head_k_max=n_sec_head_k_max,
             cond_enc=shared_cond_enc,
         )
@@ -123,6 +124,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 k_max=k_max,
                 router=stage2_router,
                 trunk_type=s2_spec.trunk.type,
+                block_conditioning=s2_spec.trunk.block_conditioning,
                 build_n_sec_head=n_sec_owner != "stage1",
                 particle_type_cfg=particle_type_cfg,
                 history=ar_cfg.history,
@@ -151,6 +153,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 k_max=k_max,
                 router=stage2_router,
                 trunk_type=s2_spec.trunk.type,
+                block_conditioning=s2_spec.trunk.block_conditioning,
                 build_n_sec_head=n_sec_owner != "stage1",
                 particle_type_cfg=particle_type_cfg,
                 cond_enc=shared_cond_enc,

giant/model/layers.py

@@ -57,6 +57,26 @@ class ContextAdapter(nn.Module):
         return torch.tanh(self.proj(x))
 
 
+BLOCK_REGISTRY: dict[str, type[nn.Module]] = {}
+
+
+def register_block(name: str):
+    def decorator(cls: type[nn.Module]) -> type[nn.Module]:
+        BLOCK_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def build_block(name: str, dim: int, cond_dim: int, dropout: float = 0.0) -> nn.Module:
+    """Factory: look up a registered conditioning-injection block by name and
+    construct one instance — `trunk.block_conditioning` (gitea #34)."""
+    if name not in BLOCK_REGISTRY:
+        raise ValueError(f"unknown block conditioning type {name!r}; available: {sorted(BLOCK_REGISTRY)}")
+    return BLOCK_REGISTRY[name](dim, cond_dim, dropout)
+
+
+@register_block("add")
 class ResBlock(nn.Module):
     def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None:
         super().__init__()
@@ -74,3 +94,66 @@ class ResBlock(nn.Module):
         h = self.dropout(h)
         h = self.linear2(h)
         return x + h
+
+
+@register_block("film")
+class FilmResBlock(nn.Module):
+    """FiLM conditioning (Perez et al. 2018): a per-channel scale+shift
+    modulates the normalized features, on top of the norm's own affine —
+    an *additional* modulation, unlike `AdaLNResBlock` below, which replaces
+    the norm's affine outright. `film_proj` is zero-initialized so
+    `gamma=beta=0` at construction — conditioning has no effect on the
+    output until training moves it, a stable starting point (though not a
+    literal identity block, since `linear1`/`linear2` aren't zero-init)."""
+
+    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.linear1 = nn.Linear(dim, dim)
+        self.film_proj = nn.Linear(cond_dim, 2 * dim)
+        nn.init.zeros_(self.film_proj.weight)
+        nn.init.zeros_(self.film_proj.bias)
+        self.act = nn.SiLU()
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        h = self.norm(x)
+        gamma, beta = self.film_proj(cond).chunk(2, dim=-1)
+        h = h * (1 + gamma) + beta
+        h = self.linear1(h)
+        h = self.act(h)
+        h = self.dropout(h)
+        h = self.linear2(h)
+        return x + h
+
+
+@register_block("adaln")
+class AdaLNResBlock(nn.Module):
+    """AdaLN-Zero conditioning (DiT, Peebles & Xie 2022): the norm's own
+    affine is replaced by a conditioning-derived scale/shift, and the
+    residual branch is scaled by a conditioning-derived gate. `adaln_proj`
+    is zero-initialized, so `scale=shift=gate=0` at construction — the block
+    is the exact identity function at init (`x + 0 * h' == x`), regardless
+    of `x`/`cond`."""
+
+    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False)
+        self.linear1 = nn.Linear(dim, dim)
+        self.adaln_proj = nn.Linear(cond_dim, 3 * dim)
+        nn.init.zeros_(self.adaln_proj.weight)
+        nn.init.zeros_(self.adaln_proj.bias)
+        self.act = nn.SiLU()
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        h = self.norm(x)
+        scale, shift, gate = self.adaln_proj(cond).chunk(3, dim=-1)
+        h = h * (1 + scale) + shift
+        h = self.linear1(h)
+        h = self.act(h)
+        h = self.dropout(h)
+        h = self.linear2(h)
+        return x + gate * h

giant/model/models.py

@@ -93,6 +93,7 @@ class Stage1Model(nn.Module):
         noise_dim: int = 64,
         router: Router | None = None,
         trunk_type: str = "resmlp",
+        block_conditioning: str = "add",
         n_sec_head_k_max: int | None = None,
         cond_enc: ConditionEncoder | None = None,
     ) -> None:
@@ -108,7 +109,9 @@ class Stage1Model(nn.Module):
         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 = noise_dim if generator == "wgan" else x_dim
-        self.trunk = build_trunk(router, trunk_type, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout)
+        self.trunk = build_trunk(
+            router, trunk_type, in_dim, x_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:
             self.n_sec_head = nn.Sequential(
@@ -187,6 +190,7 @@ class Stage2OneShot(nn.Module):
         k_max: int = K_MAX,
         router: Router | None = None,
         trunk_type: str = "resmlp",
+        block_conditioning: str = "add",
         build_n_sec_head: bool = True,
         particle_type_cfg: dict | None = None,
         cond_enc: ConditionEncoder | None = None,
@@ -214,7 +218,15 @@ class Stage2OneShot(nn.Module):
         merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
         in_dim = noise_dim if generator == "wgan" else sec_dim
         self.trunk = build_trunk(
-            router, trunk_type, in_dim, sec_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout
+            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:
@@ -333,6 +345,7 @@ class Stage2Autoregressive(nn.Module):
         k_max: int = K_MAX,
         router: Router | None = None,
         trunk_type: str = "resmlp",
+        block_conditioning: str = "add",
         build_n_sec_head: bool = True,
         particle_type_cfg: dict | None = None,
         history: str = "markov",
@@ -392,6 +405,7 @@ class Stage2Autoregressive(nn.Module):
             n_res_blocks,
             merged_cond_dim,
             dropout,
+            block_conditioning,
         )
 
         self.n_sec_head = None

giant/model/network.py

@@ -11,7 +11,17 @@ from giant.model._legacy import _migrate_legacy_model_config, migrate_legacy_sta
 from giant.model.builders import build_critics, build_models
 from giant.model.encoders import ConditionEncoder, cat_col_layout
 from giant.model.history import AttentionHistory, HistoryEncoder, MarkovHistory, _CausalAttnBlock
-from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding, _make_axis_mlp
+from giant.model.layers import (
+    BLOCK_REGISTRY,
+    AdaLNResBlock,
+    ContextAdapter,
+    FilmResBlock,
+    ResBlock,
+    SinusoidalEmbedding,
+    _make_axis_mlp,
+    build_block,
+    register_block,
+)
 from giant.model.models import (
     CriticModel,
     Stage1Model,
@@ -47,13 +57,16 @@ from giant.model.trunks import (
 )
 
 __all__ = [
+    "AdaLNResBlock",
     "AttentionHistory",
+    "BLOCK_REGISTRY",
     "ComposedRouter",
     "ConditionEncoder",
     "ContextAdapter",
     "CriticModel",
     "EnergyRouter",
     "ExpertTrunk",
+    "FilmResBlock",
     "HistoryEncoder",
     "MarkovHistory",
     "PdgRouter",
@@ -75,6 +88,7 @@ __all__ = [
     "_migrate_legacy_model_config",
     "_parse_composed_axes",
     "_route_forward",
+    "build_block",
     "build_composed_router",
     "build_critics",
     "build_expert_body",
@@ -83,6 +97,7 @@ __all__ = [
     "build_trunk",
     "cat_col_layout",
     "migrate_legacy_state_dict",
+    "register_block",
     "register_router",
     "register_trunk",
     "resolve_type_n_classes",

giant/model/trunks.py

@@ -12,7 +12,7 @@ free — no separate "routed transformer trunk" class needed.
 import torch
 import torch.nn as nn
 
-from giant.model.layers import ResBlock
+from giant.model.layers import build_block
 from giant.model.routers import Router
 
 TRUNK_REGISTRY: dict[str, type[nn.Module]] = {}
@@ -34,14 +34,19 @@ def build_expert_body(
     n_blocks: int,
     cond_dim: int,
     dropout: float = 0.0,
+    block_conditioning: str = "add",
 ) -> nn.Module:
     """Factory: look up a registered trunk body by name and construct one
     instance of it — used both for a standalone (unrouted) trunk and for each
-    expert inside a `RoutedTrunk`."""
+    expert inside a `RoutedTrunk`. `block_conditioning` selects the
+    `BLOCK_REGISTRY` entry each body's internal `ResBlock`-family blocks use
+    (`trunk.block_conditioning`, gitea #34) — an optional trailing kwarg a
+    future non-`ResBlock`-based body can simply ignore, same idiom as
+    `Trunk.forward`'s accept-and-ignore `cond_cont`/`cond_cat`."""
     if name not in TRUNK_REGISTRY:
         raise ValueError(f"unknown trunk type {name!r}; available: {sorted(TRUNK_REGISTRY)}")
     cls = TRUNK_REGISTRY[name]
-    return cls(in_dim, out_dim, hidden_dim, n_blocks, cond_dim, dropout)
+    return cls(in_dim, out_dim, hidden_dim, n_blocks, cond_dim, dropout, block_conditioning=block_conditioning)
 
 
 @register_trunk("resmlp")
@@ -65,11 +70,14 @@ class ExpertTrunk(nn.Module):
         n_blocks: int,
         cond_dim: int,
         dropout: float = 0.0,
+        block_conditioning: str = "add",
     ) -> None:
         super().__init__()
         self.out_dim = out_dim
         self.input_proj = nn.Linear(in_dim, hidden_dim)
-        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_blocks)])
+        self.blocks = nn.ModuleList(
+            [build_block(block_conditioning, hidden_dim, cond_dim, dropout) for _ in range(n_blocks)]
+        )
         self.out_proj = nn.Linear(hidden_dim, out_dim)
 
     def forward(
@@ -145,12 +153,22 @@ class RoutedTrunk(Trunk):
         n_res_blocks: int,
         cond_dim: int,
         dropout: float = 0.0,
+        block_conditioning: str = "add",
     ) -> None:
         super().__init__()
         self.router = router
         self.experts = nn.ModuleList(
             [
-                build_expert_body(trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
+                build_expert_body(
+                    trunk_type,
+                    in_dim,
+                    out_dim,
+                    hidden_dim,
+                    n_res_blocks,
+                    cond_dim,
+                    dropout,
+                    block_conditioning=block_conditioning,
+                )
                 for _ in range(router.n_experts)
             ]
         )
@@ -174,6 +192,7 @@ def build_trunk(
     n_res_blocks: int,
     cond_dim: int,
     dropout: float = 0.0,
+    block_conditioning: str = "add",
 ) -> nn.Module:
     """Build a stage's trunk: `trunk_type` (a `TRUNK_REGISTRY` key, e.g.
     `"resmlp"`) selects the expert body architecture; `router`, if given,
@@ -182,8 +201,14 @@ def build_trunk(
     class), which is what makes an unrouted trunk's state-dict keys land
     directly under `trunk.*` instead of `trunk.experts.0.*` (see
     `giant.model._legacy.migrate_legacy_state_dict`, which assumes exactly
-    this flat layout for a v0.2 monolithic checkpoint).
+    this flat layout for a v0.2 monolithic checkpoint). `block_conditioning`
+    (a `BLOCK_REGISTRY` key, e.g. `"add"`/`"film"`/`"adaln"`) selects each
+    body's conditioning-injection mechanism (gitea #34).
     """
     if router is not None:
-        return RoutedTrunk(router, trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
-    return build_expert_body(trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
+        return RoutedTrunk(
+            router, trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout, block_conditioning
+        )
+    return build_expert_body(
+        trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout, block_conditioning
+    )

tests/test_config.py

@@ -85,6 +85,15 @@ def test_trunk_config_defaults_to_resmlp_for_both_stages():
     assert gconfig.DEFAULT_CONFIG["stage2_model"]["trunk"]["type"] == "resmlp"
 
 
+def test_trunk_config_defaults_block_conditioning_to_add_for_both_stages():
+    """gitea #34: a pre-existing config with no `block_conditioning` key
+    must reproduce today's additive-bias behaviour exactly."""
+    assert gconfig.Stage1ModelConfig().trunk.block_conditioning == "add"
+    assert gconfig.Stage2ModelConfig().trunk.block_conditioning == "add"
+    assert gconfig.DEFAULT_CONFIG["stage1_model"]["trunk"]["block_conditioning"] == "add"
+    assert gconfig.DEFAULT_CONFIG["stage2_model"]["trunk"]["block_conditioning"] == "add"
+
+
 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
@@ -936,6 +945,26 @@ def test_validate_config_keys_rejects_unknown_trunk_key():
         assert "type" in str(e)
 
 
+def test_validate_config_keys_allows_block_conditioning():
+    cfg = _cfg_with(
+        **{
+            "stage1_model.trunk.block_conditioning": "film",
+            "stage2_model.trunk.block_conditioning": "adaln",
+        }
+    )
+    gconfig.validate_config_keys(cfg)  # must not raise
+
+
+def test_validate_config_keys_rejects_unknown_block_conditioning_key():
+    cfg = _cfg_with(**{"stage1_model.trunk.block_conditioning_o": "film"})  # typo
+    try:
+        gconfig.validate_config_keys(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "stage1_model.trunk.block_conditioning_o" in str(e)
+        assert "block_conditioning" in str(e)
+
+
 def test_merge_cli_overrides_rejects_typo_in_toml_file(tmp_path):
     path = tmp_path / "config.toml"
     path.write_text("[meta]\nconfig_version = 3\n\n[stage1_model]\nn_res_block = 12\n")

tests/test_router.py

@@ -5,16 +5,21 @@ import torch
 
 from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
 from giant.model.network import (
+    BLOCK_REGISTRY,
     TRUNK_REGISTRY,
+    AdaLNResBlock,
     ComposedRouter,
     EnergyRouter,
     ExpertTrunk,
+    FilmResBlock,
     PdgRouter,
     ProcessRouter,
     ROUTER_REGISTRY,
+    ResBlock,
     RoutedTrunk,
     Stage1Model,
     Stage2OneShot,
+    build_block,
     build_composed_router,
     build_expert_body,
     build_models,
@@ -162,6 +167,52 @@ def test_build_expert_body_unknown_type_raises():
     raise AssertionError("expected ValueError for unknown trunk type")
 
 
+# ── BLOCK_REGISTRY / build_block (gitea #34) ────────────────────────────────
+
+
+def test_block_registry_has_add_film_adaln():
+    assert BLOCK_REGISTRY["add"] is ResBlock
+    assert BLOCK_REGISTRY["film"] is FilmResBlock
+    assert BLOCK_REGISTRY["adaln"] is AdaLNResBlock
+
+
+def test_build_block_unknown_type_raises():
+    try:
+        build_block("nonexistent", dim=8, cond_dim=4)
+    except ValueError:
+        return
+    raise AssertionError("expected ValueError for unknown block conditioning type")
+
+
+@pytest.mark.parametrize("block_type", ["add", "film", "adaln"])
+def test_block_forward_shape(block_type):
+    block = build_block(block_type, dim=8, cond_dim=4)
+    x = torch.randn(5, 8)
+    cond = torch.randn(5, 4)
+    out = block(x, cond)
+    assert out.shape == (5, 8)
+
+
+def test_film_res_block_output_invariant_to_cond_at_init():
+    """Zero-initialized film_proj means gamma=beta=0 at construction, so the
+    output must not depend on which cond is passed in."""
+    block = FilmResBlock(dim=8, cond_dim=4)
+    x = torch.randn(5, 8)
+    cond_a = torch.randn(5, 4)
+    cond_b = torch.randn(5, 4)
+    torch.testing.assert_close(block(x, cond_a), block(x, cond_b))
+
+
+def test_adaln_res_block_is_identity_at_init():
+    """Zero-initialized adaln_proj means scale=shift=gate=0 at construction,
+    so the block must be the exact identity function (the 'Zero' in
+    AdaLN-Zero)."""
+    block = AdaLNResBlock(dim=8, cond_dim=4)
+    x = torch.randn(5, 8)
+    cond = torch.randn(5, 4)
+    torch.testing.assert_close(block(x, cond), x)
+
+
 # ── EnergyRouter learn_width / learn_temperature ────────────────────────────
 
 
@@ -1077,6 +1128,44 @@ def test_build_models_explicit_resmlp_trunk_type_matches_default():
     assert default_params == explicit_params
 
 
+def test_build_models_explicit_add_block_conditioning_matches_default():
+    """stage1_model.trunk.block_conditioning = 'add' is the default's
+    spelled-out equivalent, not a behaviour change — gitea #34."""
+    default_cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2)
+    explicit_cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2, trunk={"type": "resmlp", "block_conditioning": "add"})
+    default_stage1 = build_models(default_cfg)["stage1"]
+    explicit_stage1 = build_models(explicit_cfg)["stage1"]
+    assert default_stage1 is not None and explicit_stage1 is not None
+    assert type(default_stage1.trunk.blocks[0]) is type(explicit_stage1.trunk.blocks[0]) is ResBlock
+    default_params = sum(p.numel() for p in default_stage1.parameters())
+    explicit_params = sum(p.numel() for p in explicit_stage1.parameters())
+    assert default_params == explicit_params
+
+
+@pytest.mark.parametrize("block_type,cls", [("film", FilmResBlock), ("adaln", AdaLNResBlock)])
+def test_build_models_selects_block_conditioning(block_type, cls):
+    cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2, trunk={"type": "resmlp", "block_conditioning": block_type})
+    stage1 = build_models(cfg)["stage1"]
+    assert stage1 is not None
+    assert isinstance(stage1.trunk, ExpertTrunk)
+    assert all(isinstance(b, cls) for b in stage1.trunk.blocks)
+
+
+def test_build_models_routed_trunk_uses_block_conditioning_for_every_expert():
+    cfg = _nested_cfg(
+        pdg_vocab=4,
+        mat_vocab=2,
+        trunk={"type": "resmlp", "block_conditioning": "film"},
+        stage1_router={"enabled": True, "type": "energy", "n_experts": 3},
+    )
+    stage1 = build_models(cfg)["stage1"]
+    assert stage1 is not None
+    assert isinstance(stage1.trunk, RoutedTrunk)
+    assert len(stage1.trunk.experts) == 3
+    for expert in stage1.trunk.experts:
+        assert all(isinstance(b, FilmResBlock) for b in expert.blocks)
+
+
 def test_build_models_routed_pair_is_drop_in_for_sample_flow():
     """Exercise the exact calling convention giant/sample.py uses."""
     from giant.sample import sample_flow, sample_secondaries