giant/giant/model/network.py

import inspect
import math
import re
from collections.abc import Sequence

import torch
import torch.nn as nn
import torch.nn.functional as F

from giant.constants import (
    COND_DIM,
    COND_DIM_BASE,
    EMB_DIM,
    K_MAX,
    MATERIAL_PHYS_DIM,
    PARTICLE_PHYS_DIM,
    SEC_DIM,
    SEC_SLOT_DIM,
    X_DIM,
)

# ---------------------------------------------------------------------------
# Building blocks (docs/v0.3.0-design.md §5.2/§5.3)
# ---------------------------------------------------------------------------


class SinusoidalEmbedding(nn.Module):
    def __init__(self, dim: int) -> None:
        super().__init__()
        assert dim % 2 == 0, "dim must be even"
        half = dim // 2
        freqs = torch.exp(
            -math.log(10000)
            * torch.arange(half, dtype=torch.float32)
            / max(half - 1, 1)
        )
        self.register_buffer("freqs", freqs)

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        t = t.reshape(-1, 1).float()
        args = t * self.freqs.unsqueeze(0)  # (B, half)
        return torch.cat([args.sin(), args.cos()], dim=-1)  # (B, dim)


def _make_axis_mlp(in_dim: int, emb_dim: int, n_layers: int) -> nn.Sequential:
    """`n_layers`-deep MLP producing an `emb_dim`-wide vector from `in_dim`
    physical properties (`conditioning.{particle,material}.n_layers`).

    `n_layers=1` (the v0.3.0 default): a single `Linear`, no hidden
    activation. `n_layers=2` reproduces v0.2's hardcoded depth exactly —
    `Linear -> SiLU -> Linear` — which is why `migrate_config` back-fills
    `n_layers=2` for migrated configs rather than the v0.3 default of 1 (see
    its docstring).
    """
    if n_layers < 1:
        raise ValueError(f"n_layers must be >= 1, got {n_layers}")
    if n_layers == 1:
        return nn.Sequential(nn.Linear(in_dim, emb_dim))
    layers: list[nn.Module] = [nn.Linear(in_dim, emb_dim), nn.SiLU()]
    for _ in range(n_layers - 2):
        layers += [nn.Linear(emb_dim, emb_dim), nn.SiLU()]
    layers.append(nn.Linear(emb_dim, emb_dim))
    return nn.Sequential(*layers)


class ConditionEncoder(nn.Module):
    """Fuses continuous conditioning with particle/material identity.

    The particle and material axes are configured independently
    (`particle_cfg`/`material_cfg`, each `{"type", "emb_dim", "n_layers"}` —
    see docs/v0.3.0-design.md §3.1) and may mix freely, e.g. material
    "physical" with particle "embedding". Three modes per axis:
      - "embedding": a learned `nn.Embedding` lookup, indexed by `cond_cat`'s
        dense training-vocab index. Memorizes the training menu.
      - "physical": an `n_layers`-deep MLP over the axis's raw physical
        properties (already present in `cond_cont[:, COND_DIM_BASE:]` — see
        giant.data.transforms.build_features), computable for any PDG code /
        material name rather than only ones seen in training.
      - "onehot": not yet implemented (v0.3.0 step 4 — the top-N map isn't
        built yet); raises `NotImplementedError` if selected.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        particle_cfg: dict,
        material_cfg: dict,
        cont_dim: int = COND_DIM,
        out_dim: int = 128,
    ) -> None:
        super().__init__()
        self.particle_cfg = dict(particle_cfg)
        self.material_cfg = dict(material_cfg)

        p_type = particle_cfg["type"]
        p_emb_dim = particle_cfg["emb_dim"]
        if p_type == "embedding":
            self.pdg_emb = nn.Embedding(pdg_vocab, p_emb_dim)
        elif p_type == "physical":
            self.particle_mlp = _make_axis_mlp(
                PARTICLE_PHYS_DIM, p_emb_dim, particle_cfg.get("n_layers", 1)
            )
        elif p_type != "onehot":
            raise ValueError(f"unknown conditioning.particle.type {p_type!r}")

        m_type = material_cfg["type"]
        m_emb_dim = material_cfg["emb_dim"]
        if m_type == "embedding":
            self.mat_emb = nn.Embedding(mat_vocab, m_emb_dim)
        elif m_type == "physical":
            self.material_mlp = _make_axis_mlp(
                MATERIAL_PHYS_DIM, m_emb_dim, material_cfg.get("n_layers", 1)
            )
        elif m_type != "onehot":
            raise ValueError(f"unknown conditioning.material.type {m_type!r}")

        in_dim = COND_DIM_BASE + p_emb_dim + m_emb_dim
        self.mlp = nn.Sequential(
            nn.Linear(in_dim, out_dim),
            nn.SiLU(),
            nn.Linear(out_dim, out_dim),
        )

    def _particle_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor):
        p_type = self.particle_cfg["type"]
        if p_type == "embedding":
            return self.pdg_emb(cond_cat[:, 0])
        if p_type == "physical":
            particle_phys = cond_cont[
                :, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM
            ]
            return self.particle_mlp(particle_phys)
        raise NotImplementedError(
            "conditioning.particle.type='onehot' needs the top-N PDG map "
            "(v0.3.0 step 4, not yet implemented)"
        )

    def _material_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor):
        m_type = self.material_cfg["type"]
        if m_type == "embedding":
            return self.mat_emb(cond_cat[:, 1])
        if m_type == "physical":
            material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :]
            return self.material_mlp(material_phys)
        raise NotImplementedError(
            "conditioning.material.type='onehot' needs the top-N material "
            "map (v0.3.0 step 4, not yet implemented)"
        )

    def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        pdg_e = self._particle_embed(cond_cont, cond_cat)
        mat_e = self._material_embed(cond_cont, cond_cat)
        x = torch.cat([cond_cont[:, :COND_DIM_BASE], pdg_e, mat_e], dim=-1)
        return self.mlp(x)


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 —
    `stage2_model.context_dim`. Was `SecondaryConditionEncoder.stage1_proj`
    (+ its `tanh`) in v0.2; pulled out as its own module in v0.3.0 since
    `SecondaryConditionEncoder` as a wrapper class disappears (design doc §5.2)."""

    def __init__(self, in_dim: int, context_dim: int) -> None:
        super().__init__()
        self.proj = nn.Linear(in_dim, context_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.tanh(self.proj(x))


class ResBlock(nn.Module):
    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.cond_proj = nn.Linear(cond_dim, dim, bias=False)
        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)
        h = self.linear1(h) + self.cond_proj(cond)
        h = self.act(h)
        h = self.dropout(h)
        h = self.linear2(h)
        return x + h


# ---------------------------------------------------------------------------
# Routers — carried over unchanged from v0.2 (docs/v0.3.0-design.md §5.3)
# ---------------------------------------------------------------------------


class Router(nn.Module):
    """Contract for a pluggable mixture-of-experts routing axis.

    Subclasses implement `gate` (soft partition-of-unity weights over
    experts, used in train mode for a fully differentiable mixture);
    `top1` and `balance_loss` have working defaults so a new routing axis
    is usually a one-method add. See `ROUTER_REGISTRY` / `build_router`.
    """

    def __init__(self, n_experts: int) -> None:
        super().__init__()
        self.n_experts = n_experts
        self.gumbel = False
        self.gumbel_tau = 1.0

    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        """(B, n_experts) soft weights, rows summing to 1."""
        raise NotImplementedError

    def combine_weights(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
    ) -> torch.Tensor:
        """(B, n_experts) train-time expert-combination weights.

        Default (`gumbel=False`): identical to `gate()`. Opt-in
        straight-through Gumbel-softmax (`gumbel=True`, train mode only):
        hardens the forward pass to a one-hot sample (matching eval-time
        top-1 dispatch) while keeping the soft sample's gradient on backward.
        """
        probs = self.gate(cond_cont, cond_cat)
        if not (self.gumbel and self.training):
            return probs
        log_probs = torch.log(probs.clamp_min(1e-8))
        return F.gumbel_softmax(log_probs, tau=self.gumbel_tau, hard=True, dim=-1)

    def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        """(B,) hard expert index, used for eval-time grouped dispatch."""
        return self.gate(cond_cont, cond_cat).argmax(dim=-1)

    def balance_loss(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
    ) -> torch.Tensor:
        """Importance CV^2 load-balancing auxiliary loss (Shazeer et al. 2017)."""
        importance = self.gate(cond_cont, cond_cat).sum(dim=0)  # (n_experts,)
        return (importance.std() / (importance.mean() + 1e-8)) ** 2

    def classify_loss(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
    ) -> torch.Tensor:
        """Optional supervised auxiliary loss shaping the router's own belief.

        Default: none (a scalar 0). Routers gating on an unobservable
        pre-step quantity (e.g. ProcessRouter) override this.
        """
        return torch.zeros((), device=cond_cont.device)

    def entropy_loss(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
    ) -> torch.Tensor:
        """Optional auxiliary loss rewarding sharper (lower-entropy) routing."""
        norm_entropy, _ = self.gate_stats(cond_cont, cond_cat)
        return norm_entropy

    def gate_stats(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Diagnostics: `(norm_entropy, importance)` — see v0.2 docstring for
        the full explanation, unchanged in v0.3.0."""
        gate = self.gate(cond_cont, cond_cat)  # (B, n_experts)
        row_entropy = -(gate * (gate + 1e-8).log()).sum(dim=-1)  # (B,)
        norm_entropy = row_entropy.mean() / math.log(self.n_experts)
        importance = gate.sum(dim=0)  # (n_experts,)
        return norm_entropy, importance


ROUTER_REGISTRY: dict[str, type[Router]] = {}


def register_router(name: str):
    def decorator(cls: type[Router]) -> type[Router]:
        ROUTER_REGISTRY[name] = cls
        return cls

    return decorator


def build_router(name: str, n_experts: int, **kwargs) -> Router:
    """Factory: look up a `Router` subclass by name from the registry.

    Every registered router type is fed the same `router` config dict;
    kwargs not declared by that type's constructor are silently dropped, so
    per-type hyperparameters (e.g. EnergyRouter's `temperature`) can coexist
    in one config without special-casing.
    """
    if name not in ROUTER_REGISTRY:
        raise ValueError(
            f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}"
        )
    cls = ROUTER_REGISTRY[name]
    accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"}
    filtered = {k: v for k, v in kwargs.items() if k in accepted}
    return cls(n_experts=n_experts, **filtered)


def _bounded_interp(raw: torch.Tensor, lo: float, hi: float) -> torch.Tensor:
    """Sigmoid interpolation into `[lo, hi]` — smooth, always-positive-gradient
    bound used for EnergyRouter's `learn_width`/`learn_temperature` modes."""
    return lo + (hi - lo) * torch.sigmoid(raw)


def _inverse_bounded_interp(value: float, lo: float, hi: float) -> float:
    """Inverse of `_bounded_interp`, used once at construction to warm-start
    `raw` so the initial effective width/temperature exactly equals `value`."""
    p = min(max((value - lo) / (hi - lo), 1e-6), 1 - 1e-6)
    return math.log(p / (1 - p))


@register_router("energy")
class EnergyRouter(Router):
    """Soft turn-on gate over normalized pre-step log-energy.

    Reads `cond_cont[:, energy_idx]` (ignores cond_cat). `gate(e) =
    softmax_i(-(e - c_i)^2 / tau)`; as tau -> 0 this hardens to
    nearest-center (Voronoi) selection, exactly what `top1` uses at eval.
    """

    def __init__(
        self,
        n_experts: int = 4,
        temperature: float = 0.5,
        learn_centers: bool = True,
        energy_idx: int = 3,
        centers_init: Sequence[float] | None = None,
        learn_width: bool = False,
        learn_temperature: bool = False,
        width_min_ratio: float = 0.1,
        width_max_ratio: float = 10.0,
    ) -> None:
        super().__init__(n_experts)
        if learn_width and learn_temperature:
            raise ValueError("learn_width and learn_temperature are mutually exclusive")
        self.temperature = temperature
        self.energy_idx = energy_idx
        self.learn_width = learn_width
        self.learn_temperature = learn_temperature
        if learn_width or learn_temperature:
            if not (width_min_ratio < 1.0 < width_max_ratio):
                raise ValueError(
                    f"width_min_ratio ({width_min_ratio}) and width_max_ratio "
                    f"({width_max_ratio}) must bracket 1.0"
                )
            self._width_lo = width_min_ratio * temperature
            self._width_hi = width_max_ratio * temperature
            raw0 = _inverse_bounded_interp(temperature, self._width_lo, self._width_hi)
            if learn_width:
                self.raw_width = nn.Parameter(torch.full((n_experts,), raw0))
            else:
                self.raw_temperature = nn.Parameter(torch.tensor(raw0))
        if centers_init is None:
            centers = torch.linspace(-2.0, 2.0, n_experts)
        else:
            if len(centers_init) != n_experts:
                raise ValueError(
                    f"centers_init has {len(centers_init)} values, "
                    f"expected n_experts={n_experts}"
                )
            centers = torch.tensor(list(centers_init), dtype=torch.float32)
        if learn_centers:
            self.centers = nn.Parameter(centers)
        else:
            self.register_buffer("centers", centers)

    def effective_width(self) -> torch.Tensor | float:
        if self.learn_width:
            return _bounded_interp(self.raw_width, self._width_lo, self._width_hi)
        if self.learn_temperature:
            return _bounded_interp(self.raw_temperature, self._width_lo, self._width_hi)
        return self.temperature

    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        e = cond_cont[:, self.energy_idx].unsqueeze(-1)  # (B, 1)
        d2 = (e - self.centers.unsqueeze(0)) ** 2  # (B, n_experts)
        return torch.softmax(-d2 / self.effective_width(), dim=-1)


@register_router("pdg")
class PdgRouter(Router):
    """Soft turn-on gate over a learned PDG embedding (own table, separate
    from the trunk's `ConditionEncoder`). No supervision needed — PDG code
    is already known at pre-step time."""

    def __init__(
        self,
        n_experts: int,
        pdg_vocab: int,
        emb_dim: int = 8,
        temperature: float = 0.5,
        learn_centers: bool = True,
    ) -> None:
        super().__init__(n_experts)
        self.temperature = temperature
        self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
        centers = torch.randn(n_experts, emb_dim) * 0.1
        if learn_centers:
            self.centers = nn.Parameter(centers)
        else:
            self.register_buffer("centers", centers)

    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        e = self.pdg_emb(cond_cat[:, 0])  # (B, emb_dim)
        d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(
            -1
        )  # (B, n_experts)
        return torch.softmax(-d2 / self.temperature, dim=-1)


@register_router("process")
class ProcessRouter(Router):
    """Routes on the physics process expected to end the step — a post-step
    outcome, so a small classifier over pre-step conditioning predicts it
    (own pdg/material embeddings, separate from the trunk's ConditionEncoder).
    `n_experts` doubles as the number of process classes. Supervised via
    `classify_loss` against the true `process` label at train time only;
    `gate`/`top1` never see it."""

    def __init__(
        self,
        n_experts: int,
        pdg_vocab: int,
        mat_vocab: int,
        emb_dim: int = 8,
        hidden_dim: int = 64,
    ) -> None:
        super().__init__(n_experts)
        self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
        self.mat_emb = nn.Embedding(mat_vocab, emb_dim)
        self.classifier = nn.Sequential(
            nn.Linear(COND_DIM + 2 * emb_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, n_experts),
        )

    def logits(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        pdg_e = self.pdg_emb(cond_cat[:, 0])
        mat_e = self.mat_emb(cond_cat[:, 1])
        h = torch.cat([cond_cont, pdg_e, mat_e], dim=-1)
        return self.classifier(h)

    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        return torch.softmax(self.logits(cond_cont, cond_cat), dim=-1)

    def classify_loss(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
    ) -> torch.Tensor:
        return F.cross_entropy(self.logits(cond_cont, cond_cat), labels)


class ComposedRouter(Router):
    """Joint router over independent axes (e.g. energy x pdg), outer-product
    gated. Not registered in `ROUTER_REGISTRY`; use `build_composed_router`."""

    def __init__(self, routers: list[Router]) -> None:
        if not routers:
            raise ValueError("ComposedRouter needs at least one sub-router")
        n_experts = 1
        for r in routers:
            n_experts *= r.n_experts
        super().__init__(n_experts)
        self.routers = nn.ModuleList(routers)

    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        joint = self.routers[0].gate(cond_cont, cond_cat)  # (B, n_0)
        for router in self.routers[1:]:
            g = router.gate(cond_cont, cond_cat)  # (B, n_i)
            joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(
                1
            )  # (B, prod so far)
        return joint

    def classify_loss(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
    ) -> torch.Tensor:
        total = torch.zeros((), device=cond_cont.device)
        for router in self.routers:
            total = total + router.classify_loss(cond_cont, cond_cat, labels)
        return total


def build_composed_router(specs: list[dict], **shared_kwargs) -> ComposedRouter:
    """Build a `ComposedRouter` from a list of per-axis router specs — see
    `_parse_composed_axes`."""
    routers = [
        build_router(
            spec["type"],
            spec["n_experts"],
            **{
                **shared_kwargs,
                **{k: v for k, v in spec.items() if k not in ("type", "n_experts")},
            },
        )
        for spec in specs
    ]
    return ComposedRouter(routers)


_AXIS_KEY_RE = re.compile(r"^axis(\d+)_(.+)$")


def _parse_composed_axes(router_cfg: dict) -> list[dict]:
    """Regroup `axis{i}_{field}` flat keys into a list of per-axis spec dicts.

    e.g. `axis0_type = "energy"`, `axis0_n_experts = 4`, `axis1_type = "pdg"`,
    `axis1_n_experts = 3`, `axis1_emb_dim = 8`. Axis indices must be
    contiguous from 0.
    """
    axes: dict[int, dict] = {}
    for key, value in router_cfg.items():
        m = _AXIS_KEY_RE.match(key)
        if m is None:
            continue
        idx, field = int(m.group(1)), m.group(2)
        axes.setdefault(idx, {})[field] = value
    missing = set(range(len(axes))) - axes.keys()
    if missing:
        raise ValueError(f"composed router config has gaps at axis indices {missing}")
    return [axes[i] for i in range(len(axes))]


# Router types that read cond_cat's pdg index through their own
# nn.Embedding(pdg_vocab, ...), regardless of the trunk's particle
# conditioning mode — see _check_router_conditioning_compat.
_VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process")


def _check_router_conditioning_compat(
    router_types: list[str], particle_conditioning: str
) -> None:
    """Reject a router axis that reintroduces a training-vocab PDG lookup
    under `conditioning.particle.type = "physical"`.

    `PdgRouter`/`ProcessRouter` always build their own dataset-scoped
    `nn.Embedding(pdg_vocab, ...)`, independent of `ConditionEncoder`'s
    particle mode. Pairing either with `"physical"` would silently
    reintroduce a training-menu-scoped lookup at the routing layer,
    defeating the point of physical-property conditioning. Raised loudly at
    model-build time.
    """
    bad = sorted(set(router_types) & set(_VOCAB_SCOPED_ROUTER_TYPES))
    if bad and particle_conditioning == "physical":
        raise ValueError(
            f"router type(s) {bad} always use a training-vocab PDG embedding, "
            "which is incompatible with conditioning.particle.type='physical' "
            "(whose whole point is generalizing beyond that vocab) — pick a "
            "different router type (e.g. 'energy') or use "
            "conditioning.particle.type='embedding'."
        )


def _build_router_from_cfg(
    router_cfg: dict,
    pdg_vocab: int,
    mat_vocab: int,
    particle_conditioning: str = "embedding",
) -> Router:
    """Resolve one stage's `router` config into a `Router`, single-axis or
    composed. `gumbel` is set as a post-construction attribute (shared by
    every router type, not a per-type constructor kwarg)."""
    shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
    if router_cfg["type"] == "composed":
        axes = _parse_composed_axes(router_cfg)
        _check_router_conditioning_compat(
            [a["type"] for a in axes], particle_conditioning
        )
        router = build_composed_router(axes, **shared_vocab)
        router.gumbel = bool(router_cfg.get("gumbel", False))
        return router
    _check_router_conditioning_compat([router_cfg["type"]], particle_conditioning)
    router_kwargs = {
        k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")
    }
    router_kwargs.setdefault("pdg_vocab", pdg_vocab)
    router_kwargs.setdefault("mat_vocab", mat_vocab)
    router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
    router.gumbel = bool(router_cfg.get("gumbel", False))
    return router


# ---------------------------------------------------------------------------
# Trunks (docs/v0.3.0-design.md §5.2 (b))
# ---------------------------------------------------------------------------


class ExpertTrunk(nn.Module):
    """One small expert: `input_proj -> ResBlock stack -> out_proj`.

    Unlike v0.2, `out_dim` is independent of `in_dim` — needed by stage-2 AR
    tokens later (`noise_dim` in, `4 + type_dim` out), even though every
    step-2/3 caller still has `in_dim == out_dim`.
    """

    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        hidden_dim: int,
        n_blocks: int,
        cond_dim: int,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        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.out_proj = nn.Linear(hidden_dim, out_dim)

    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        x = self.input_proj(x)
        for block in self.blocks:
            x = block(x, cond)
        return self.out_proj(x)


def _route_forward(
    experts: nn.ModuleList,
    router: Router,
    x: torch.Tensor,
    cond: torch.Tensor,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    training: bool,
) -> torch.Tensor:
    """Shared dispatch for `RoutedTrunk`.

    Train mode: full mixture `sum_i weight_i * expert_i(x)` — always
    N-expert dense compute, fully differentiable (`weight` is
    `router.combine_weights`). Eval mode: grouped top-1 dispatch — each row
    runs exactly one expert, the actual source of the per-call speedup.
    """
    if training:
        weights = router.combine_weights(cond_cont, cond_cat)  # (B, n_experts)
        out = torch.zeros(x.shape[0], experts[0].out_proj.out_features, device=x.device)
        for i, expert in enumerate(experts):
            out = out + weights[:, i : i + 1] * expert(x, cond)
        return out

    idx = router.top1(cond_cont, cond_cat)  # (B,)
    out_dim = experts[0].out_proj.out_features
    out = torch.zeros(x.shape[0], out_dim, device=x.device)
    for i, expert in enumerate(experts):
        mask = idx == i
        if mask.any():
            out[mask] = expert(x[mask], cond[mask])
    return out


class Trunk(nn.Module):
    """Interface implemented by `MonolithicTrunk`/`RoutedTrunk`: everything
    downstream of the fused conditioning vector, i.e. the actual generative
    trunk of a stage (`input_proj -> blocks -> out_proj`, monolithic or
    expert-routed)."""

    def forward(
        self,
        x: torch.Tensor,
        cond: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        raise NotImplementedError


class MonolithicTrunk(Trunk):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        hidden_dim: int,
        n_res_blocks: int,
        cond_dim: int,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        self.input_proj = nn.Linear(in_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, cond_dim, dropout=dropout)
                for _ in range(n_res_blocks)
            ]
        )
        self.out_proj = nn.Linear(hidden_dim, out_dim)

    def forward(
        self,
        x: torch.Tensor,
        cond: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        x = self.input_proj(x)
        for block in self.blocks:
            x = block(x, cond)
        return self.out_proj(x)


class RoutedTrunk(Trunk):
    def __init__(
        self,
        router: Router,
        in_dim: int,
        out_dim: int,
        hidden_dim: int,
        n_res_blocks: int,
        cond_dim: int,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        self.router = router
        self.experts = nn.ModuleList(
            [
                ExpertTrunk(
                    in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout
                )
                for _ in range(router.n_experts)
            ]
        )

    def forward(
        self,
        x: torch.Tensor,
        cond: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        return _route_forward(
            self.experts, self.router, x, cond, cond_cont, cond_cat, self.training
        )


def build_trunk(
    router: Router | None,
    in_dim: int,
    out_dim: int,
    hidden_dim: int,
    n_res_blocks: int,
    cond_dim: int,
    dropout: float = 0.0,
) -> Trunk:
    if router is not None:
        return RoutedTrunk(
            router, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout
        )
    return MonolithicTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)


# ---------------------------------------------------------------------------
# Stage models (docs/v0.3.0-design.md §5.3)
# ---------------------------------------------------------------------------


class Stage1Model(nn.Module):
    """Predicts the 9D primary post-step vector. No `n_sec_head` — decision 1
    (docs/v0.3.0-design.md §2) moves 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,
        mat_vocab: int,
        particle_cfg: dict,
        material_cfg: dict,
        hidden_dim: int = 256,
        n_res_blocks: int = 6,
        cond_out_dim: int = 128,
        x_dim: int = X_DIM,
        dropout: float = 0.0,
        generator: str = "flow",
        time_dim: int = 64,
        noise_dim: int = 64,
        router: Router | None = None,
        n_sec_head_k_max: int | None = None,
    ) -> None:
        super().__init__()
        self.generator_kind = generator
        self.noise_dim = noise_dim
        self.cond_enc = ConditionEncoder(
            pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
        )
        has_time = generator in ("flow", "ddpm")
        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, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout
        )
        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),
            )

    def forward(
        self,
        x_t: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        t: torch.Tensor | None = None,
    ) -> torch.Tensor:
        c_emb = self.cond_enc(cond_cont, cond_cat)
        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 (decision 1)."""
        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 (decision 1); this method only exists "
                "for a migrated v0.2 checkpoint (legacy_owner='stage1')"
            )
        c_emb = self.cond_enc(cond_cont, cond_cat)
        return self.n_sec_head(c_emb)


class Stage2OneShot(nn.Module):
    """Predicts all `k_max` secondary slots simultaneously — v0.2 behaviour,
    reproduced exactly (see docs/v0.3.0-design.md §12 step 2 acceptance
    criterion; `decoder = "autoregressive"` is `Stage2Autoregressive`,
    step 4/5, not implemented yet).

    Owns `n_sec_head` by default (decision 1) unless `build_n_sec_head=False`
    (a migrated v0.2 checkpoint, whose n_sec_head instead attaches to
    Stage1Model — see `_migrate_legacy_model_config`).
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        particle_cfg: dict,
        material_cfg: dict,
        hidden_dim: int = 256,
        n_res_blocks: int = 6,
        cond_out_dim: int = 128,
        context_dim: int = 64,
        sec_dim: int = SEC_DIM,
        x_dim: int = X_DIM,
        dropout: float = 0.0,
        generator: str = "wgan",
        time_dim: int = 64,
        noise_dim: int = 64,
        k_max: int = K_MAX,
        router: Router | None = None,
        build_n_sec_head: bool = True,
    ) -> None:
        super().__init__()
        self.generator_kind = generator
        self.noise_dim = noise_dim
        self.cond_enc = ConditionEncoder(
            pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
        )
        self.context_adapter = ContextAdapter(x_dim, context_dim)
        self.fuse = nn.Sequential(
            nn.Linear(cond_out_dim + context_dim, cond_out_dim),
            nn.SiLU(),
        )
        has_time = generator in ("flow", "ddpm")
        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 sec_dim
        self.trunk = build_trunk(
            router, in_dim, sec_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout
        )
        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),
            )

    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)
        ctx = self.context_adapter(stage1_out)
        return self.fuse(torch.cat([base, ctx], dim=-1))

    def forward(
        self,
        x_t: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor,
        t: torch.Tensor | None = None,
    ) -> torch.Tensor:
        c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
        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,
        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 (legacy_owner='stage1'); call "
                "stage1.predict_n_sec(cond_cont, cond_cat) instead"
            )
        c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
        return self.n_sec_head(c_emb)


class Stage2Autoregressive(nn.Module):
    """Not implemented until v0.3.0 steps 4-7 (docs/v0.3.0-design.md §6, §12)
    — stub so `stage2_model.decoder = "autoregressive"` (the v0.3.0 default)
    fails loudly instead of silently no-op-ing."""

    def __init__(self, *args, **kwargs) -> None:
        super().__init__()
        raise NotImplementedError(
            "stage2_model.decoder = 'autoregressive' is not implemented yet "
            "(design doc v0.3.0 steps 4-7) — use decoder = 'one_shot' for now"
        )


class CriticModel(nn.Module):
    """Generator-agnostic WGAN-GP critic body: a scalar realism score, for
    either stage (`stage="stage1"` mirrors v0.2 `Critic`; `stage="stage2"`
    mirrors v0.2 `SecondaryCritic`, adding the same context-fusion path as
    `Stage2OneShot`). Used only when that stage's `generator == "wgan"`."""

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        particle_cfg: dict,
        material_cfg: dict,
        in_dim: int,
        hidden_dim: int = 256,
        n_res_blocks: int = 6,
        cond_out_dim: int = 128,
        dropout: float = 0.0,
        stage: str = "stage1",
        context_dim: int = 64,
        context_in_dim: int = X_DIM,
    ) -> None:
        super().__init__()
        if stage not in ("stage1", "stage2"):
            raise ValueError(f"stage must be 'stage1' or 'stage2', got {stage!r}")
        self.stage = stage
        self.cond_enc = ConditionEncoder(
            pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim
        )
        if stage == "stage2":
            self.context_adapter = ContextAdapter(context_in_dim, context_dim)
            self.fuse = nn.Sequential(
                nn.Linear(cond_out_dim + context_dim, cond_out_dim),
                nn.SiLU(),
            )
        self.input_proj = nn.Linear(in_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
                for _ in range(n_res_blocks)
            ]
        )
        self.out_norm = nn.LayerNorm(hidden_dim)
        self.out_proj = nn.Linear(hidden_dim, 1)

    def forward(
        self,
        x: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor | None = None,
    ) -> torch.Tensor:
        base = self.cond_enc(cond_cont, cond_cat)
        if self.stage == "stage2":
            ctx = self.context_adapter(stage1_out)
            cond = self.fuse(torch.cat([base, ctx], dim=-1))
        else:
            cond = base
        h = self.input_proj(x)
        for block in self.blocks:
            h = block(h, cond)
        return self.out_proj(self.out_norm(h)).squeeze(-1)


# ---------------------------------------------------------------------------
# v0.2 -> v0.3 checkpoint migration (docs/v0.3.0-design.md §4.1, §4.3)
# ---------------------------------------------------------------------------


def _migrate_legacy_model_config(model_config: dict) -> dict:
    """Translate a v0.2 checkpoint's flat `model_config` (giant/pipeline.py's
    old shape: `hidden_dim`/`n_blocks`/`emb_dim`/`dropout`/`conditioning`/
    `router`/`mode`/... all at one level) into the nested
    `{"pdg_vocab", "mat_vocab", "conditioning", "stage1_model",
    "stage2_model"}` shape `build_models` expects.

    Sets `stage2_model.n_sec.legacy_owner = "stage1"` so the n_sec_head
    weights a v0.2 checkpoint carries on its Stage-1 module keep loading
    there (design doc §4.1) instead of the new default location
    (`Stage2OneShot`) — the n_sec head was trained against Stage 1's own
    `ConditionEncoder` output, so it has to stay attached to Stage 1's
    module, not just be labeled as such.

    Only the monolithic (non-routed) trunk shape is exercised by the step-2
    migration test (docs/v0.3.0-design.md §4.3); a routed v0.2 checkpoint
    still builds correctly here (the router config passes through), but its
    state dict isn't covered by `migrate_legacy_state_dict` below.
    """
    m = model_config
    conditioning_mode = m.get("conditioning", "embedding")
    generator = m.get("mode", "flow")
    hidden_dim = m.get("hidden_dim", 256)
    n_blocks = m.get("n_blocks", 6)
    emb_dim = m.get("emb_dim", EMB_DIM)
    dropout = m.get("dropout", 0.1)
    k_max = m.get("k_max", K_MAX)
    noise_dim = m.get("noise_dim", 64)
    router_cfg = dict(m.get("router") or {})
    router_cfg.setdefault("enabled", False)

    return {
        "pdg_vocab": m["pdg_vocab"],
        "mat_vocab": m["mat_vocab"],
        "conditioning": {
            "out_dim": 128,
            "share_stages": False,
            "particle": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": 2},
            "material": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": 2},
        },
        "stage1_model": {
            "active": True,
            "generator": generator,
            "hidden_dim": hidden_dim,
            "n_res_blocks": n_blocks,
            "dropout": dropout,
            "flow": {"time_dim": 64},
            "ddpm": {"time_dim": 64},
            "wgan": {"noise_dim": noise_dim},
            "router": dict(router_cfg),
        },
        "stage2_model": {
            "active": True,
            "decoder": "one_shot",
            "generator": generator,
            "hidden_dim": hidden_dim,
            "n_res_blocks": n_blocks,
            "dropout": dropout,
            "k_max": k_max,
            "context_dim": 64,
            "n_sec": {"mode": "head", "legacy_owner": "stage1"},
            "particle_type": {"target": "physical"},
            "flow": {"time_dim": 64},
            "ddpm": {"time_dim": 64},
            "wgan": {"noise_dim": noise_dim},
            "router": {**router_cfg, "tie_to_stage1": False},
        },
    }


def migrate_legacy_state_dict(
    old_stage1_sd: dict, old_stage2_sd: dict
) -> tuple[dict, dict]:
    """Remap a v0.2 checkpoint's (`DenoisingMLP`-or-`WGANGenerator`,
    `SecondaryDecoder`-or-`WGANSecondaryGenerator`) state dicts onto the new
    `(Stage1Model, Stage2OneShot)` module structure produced by
    `build_models(_migrate_legacy_model_config(model_config))`.

    Only the monolithic (non-routed) trunk shape is handled — see
    docs/v0.3.0-design.md §4.3's migration test scope.
    """

    def _trunk_prefix(k: str) -> str:
        if k.startswith(("input_proj.", "blocks.", "out_proj.")):
            return f"trunk.{k}"
        return k

    new_stage1 = {}
    for k, v in old_stage1_sd.items():
        if k.startswith("n_sec_head."):
            new_stage1[k] = v  # stays top-level (legacy_owner="stage1")
        else:
            new_stage1[_trunk_prefix(k)] = v

    new_stage2 = {}
    for k, v in old_stage2_sd.items():
        if k.startswith("cond_enc.base."):
            new_stage2["cond_enc." + k[len("cond_enc.base.") :]] = v
        elif k.startswith("cond_enc.stage1_proj."):
            new_stage2["context_adapter.proj." + k[len("cond_enc.stage1_proj.") :]] = v
        elif k.startswith("cond_enc.fuse."):
            new_stage2["fuse." + k[len("cond_enc.fuse.") :]] = v
        else:
            new_stage2[_trunk_prefix(k)] = v

    return new_stage1, new_stage2


# ---------------------------------------------------------------------------
# Factories (docs/v0.3.0-design.md §5.4)
# ---------------------------------------------------------------------------


def build_models(model_config: dict) -> dict[str, nn.Module | None]:
    """Construct `{"stage1": ..., "stage2": ...}` from a config dict — either
    the new nested shape (has a `"stage1_model"` key, plus `"pdg_vocab"`/
    `"mat_vocab"`/`"conditioning"` at the top level) or a v0.2 checkpoint's
    flat `model_config`, auto-migrated via `_migrate_legacy_model_config`.

    A stage is `None` in the result when that stage's `active = False`.
    `stage2_model.router.tie_to_stage1` shares stage 1's literal `Router`
    instance rather than building a second, independently-parameterized one
    (v0.2's actual — probably accidental — behaviour: two routers built from
    one config with no semantic relationship between them).
    """
    cfg = (
        model_config
        if "stage1_model" in model_config
        else _migrate_legacy_model_config(model_config)
    )
    pdg_vocab = cfg["pdg_vocab"]
    mat_vocab = cfg["mat_vocab"]
    conditioning = cfg["conditioning"]
    if conditioning.get("share_stages"):
        raise NotImplementedError(
            "conditioning.share_stages = true is not implemented yet — each "
            "stage always builds its own ConditionEncoder for now"
        )
    particle_cfg = conditioning["particle"]
    material_cfg = conditioning["material"]
    particle_conditioning = particle_cfg["type"]
    s1cfg = cfg["stage1_model"]
    s2cfg = cfg["stage2_model"]
    cond_out_dim = conditioning.get("out_dim", 128)

    result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None}

    stage1_router: Router | None = None
    if s1cfg.get("active", True):
        router_cfg = s1cfg.get("router") or {}
        if router_cfg.get("enabled"):
            stage1_router = _build_router_from_cfg(
                router_cfg, pdg_vocab, mat_vocab, particle_conditioning
            )
        generator = s1cfg.get("generator", "flow")
        gen_sub = s1cfg.get(generator, {}) or {}
        legacy_owner = (s2cfg.get("n_sec") or {}).get("legacy_owner")
        n_sec_head_k_max = (
            s2cfg.get("k_max", K_MAX) if legacy_owner == "stage1" else None
        )
        result["stage1"] = Stage1Model(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            particle_cfg=particle_cfg,
            material_cfg=material_cfg,
            hidden_dim=s1cfg.get("hidden_dim", 256),
            n_res_blocks=s1cfg.get("n_res_blocks", 6),
            cond_out_dim=cond_out_dim,
            dropout=s1cfg.get("dropout", 0.0),
            generator=generator,
            time_dim=gen_sub.get("time_dim", 64),
            noise_dim=(s1cfg.get("wgan") or {}).get("noise_dim", 64),
            router=stage1_router,
            n_sec_head_k_max=n_sec_head_k_max,
        )

    if s2cfg.get("active", True):
        decoder = s2cfg.get("decoder", "one_shot")
        if decoder == "autoregressive":
            result["stage2"] = Stage2Autoregressive()
            return result
        router_cfg = s2cfg.get("router") or {}
        stage2_router: Router | None = None
        if router_cfg.get("enabled"):
            if router_cfg.get("tie_to_stage1") and stage1_router is not None:
                stage2_router = stage1_router
            else:
                stage2_router = _build_router_from_cfg(
                    router_cfg, pdg_vocab, mat_vocab, particle_conditioning
                )
        generator = s2cfg.get("generator", "wgan")
        gen_sub = s2cfg.get(generator, {}) or {}
        legacy_owner = (s2cfg.get("n_sec") or {}).get("legacy_owner")
        k_max = s2cfg.get("k_max", K_MAX)
        result["stage2"] = Stage2OneShot(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            particle_cfg=particle_cfg,
            material_cfg=material_cfg,
            hidden_dim=s2cfg.get("hidden_dim", 256),
            n_res_blocks=s2cfg.get("n_res_blocks", 6),
            cond_out_dim=cond_out_dim,
            context_dim=s2cfg.get("context_dim", 64),
            sec_dim=k_max * SEC_SLOT_DIM,
            dropout=s2cfg.get("dropout", 0.0),
            generator=generator,
            time_dim=gen_sub.get("time_dim", 64),
            noise_dim=(s2cfg.get("wgan") or {}).get("noise_dim", 64),
            k_max=k_max,
            router=stage2_router,
            build_n_sec_head=legacy_owner != "stage1",
        )

    return result


def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
    """Construct `{"stage1": ..., "stage2": ...}` critics for `generator =
    "wgan"` training. Training-only — never persisted for inference the way
    `build_models`'s pair is. `None` for a stage that's inactive or not
    WGAN."""
    cfg = (
        model_config
        if "stage1_model" in model_config
        else _migrate_legacy_model_config(model_config)
    )
    pdg_vocab = cfg["pdg_vocab"]
    mat_vocab = cfg["mat_vocab"]
    conditioning = cfg["conditioning"]
    particle_cfg = conditioning["particle"]
    material_cfg = conditioning["material"]
    cond_out_dim = conditioning.get("out_dim", 128)
    s1cfg = cfg["stage1_model"]
    s2cfg = cfg["stage2_model"]

    result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None}

    if s1cfg.get("active", True) and s1cfg.get("generator") == "wgan":
        result["stage1"] = CriticModel(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            particle_cfg=particle_cfg,
            material_cfg=material_cfg,
            in_dim=X_DIM,
            hidden_dim=s1cfg.get("hidden_dim", 256),
            n_res_blocks=s1cfg.get("n_res_blocks", 6),
            cond_out_dim=cond_out_dim,
            dropout=s1cfg.get("dropout", 0.0),
            stage="stage1",
        )

    if (
        s2cfg.get("active", True)
        and s2cfg.get("generator") == "wgan"
        and s2cfg.get("decoder", "one_shot") != "autoregressive"
    ):
        k_max = s2cfg.get("k_max", K_MAX)
        result["stage2"] = CriticModel(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            particle_cfg=particle_cfg,
            material_cfg=material_cfg,
            in_dim=k_max * SEC_SLOT_DIM,
            hidden_dim=s2cfg.get("hidden_dim", 256),
            n_res_blocks=s2cfg.get("n_res_blocks", 6),
            cond_out_dim=cond_out_dim,
            dropout=s2cfg.get("dropout", 0.0),
            stage="stage2",
            context_dim=s2cfg.get("context_dim", 64),
        )

    return result