giant/giant/model/models.py

"""Top-level stage models: `Stage1Model`, `Stage2OneShot`, `Stage2Autoregressive`,
`CriticModel` — composed from encoders/trunks/history (issues.md Issue 8)."""

import torch
import torch.nn as nn

from giant.config import ConditioningAxisConfig, HeadConfig, ParticleTypeConfig
from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, SEC_SLOT_DIM, X_DIM
from giant.model.encoders import ConditionEncoder
from giant.model.history import HistoryEncoder, build_history
from giant.model.layers import ContextAdapter, SinusoidalEmbedding, build_mlp_head
from giant.model.objectives import build_objective
from giant.model.routers import Router
from giant.model.trunks import build_trunk

# ---------------------------------------------------------------------------
# Stage models
# ---------------------------------------------------------------------------


def resolve_type_n_classes(particle_type_cfg: ParticleTypeConfig, particle_emb_dim: int) -> int:
    """Effective width fed to `stage2_type_dim`/`stage2_trunk_sec_dim` in
    place of a bare `conditioning.particle.emb_dim` read. Under
    `target = "onehot"` this is `stage2_model.particle_type.n_classes` (0 =
    inherit `conditioning.particle.emb_dim`) — see gitea #29, which decoupled
    the secondary-species vocabulary size from the unrelated
    physical-conditioning MLP's output width. Under `target = "embedding"`
    (or `"physical"`, which ignores this value entirely) `n_classes` doesn't
    apply — the width stays `conditioning.particle.emb_dim`, the embedding
    table's own dimensionality (`validate_config` requires
    `conditioning.particle.type = "embedding"` here)."""
    if particle_type_cfg.target == "onehot":
        return particle_type_cfg.n_classes or particle_emb_dim
    return particle_emb_dim


def stage2_type_dim(particle_type_cfg: ParticleTypeConfig, emb_dim: int) -> int:
    """Width of a single secondary slot's type slice —
    `PARTICLE_PHYS_DIM` (log_mass, charge) for `target = "physical"`, else
    `emb_dim` (both `"onehot"` class logits and `"embedding"` vectors are
    this many classes/dims wide — callers resolve `emb_dim` via
    `resolve_type_n_classes` first)."""
    return PARTICLE_PHYS_DIM if particle_type_cfg.target == "physical" else emb_dim


def stage2_trunk_sec_dim(particle_type_cfg: ParticleTypeConfig, generator: str, k_max: int, emb_dim: int) -> int:
    """`Stage2OneShot`'s trunk output width.

    `target = "physical"` is untouched from v0.2/today:
    `k_max * SEC_SLOT_DIM`, the type slice folded into the same
    flow-matched/WGAN vector as the continuous stick/dir slots.

    `target` in `("onehot", "embedding")`: under an objective with
    `folds_type_slice` (currently just wgan) the type slice is still folded
    in (adversarial for onehot via ST-Gumbel, already-continuous for
    embedding), just `emb_dim` wide instead of `PARTICLE_PHYS_DIM` wide:
    `k_max * (CONT_SLOT_DIM + emb_dim)`. Otherwise (flow/ddpm) the type slice
    isn't part of this vector at all — it's `Stage2OneShot.type_head`'s job
    instead — so the trunk only covers `k_max * CONT_SLOT_DIM`.
    """
    if particle_type_cfg.target == "physical":
        return k_max * SEC_SLOT_DIM
    if build_objective(generator).folds_type_slice:
        return k_max * (CONT_SLOT_DIM + emb_dim)
    return k_max * CONT_SLOT_DIM


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: ConditioningAxisConfig,
        material_cfg: ConditioningAxisConfig,
        cond_out_dim: int,
        generator: str,
        noise_dim: int,
        k_max: int | None = None,
        particle_type_cfg: ParticleTypeConfig | None = None,
        cond_enc: ConditionEncoder | None = None,
    ) -> None:
        super().__init__()
        self.generator_kind = generator
        self.noise_dim = noise_dim
        self.k_max = k_max
        # `ParticleTypeConfig()`'s own dataclass default is target="onehot"
        # (the config.toml default when [stage2_model.particle_type] is
        # omitted) — a different question from "nobody passed anything to
        # this constructor", which direct/test construction relies on
        # defaulting to "physical" (build_models/build_critics always pass
        # particle_type_cfg explicitly, so this sentinel is never hit there).
        self.particle_type_cfg = (
            particle_type_cfg if particle_type_cfg is not None else ParticleTypeConfig(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,
        build_stop_head: bool = False,
        stop_head_cfg: dict | None = None,
    ) -> None:
        """Builds `self.time_emb`, `self.trunk`, `self.n_sec_head`,
        `self.type_head`, `self.stop_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). `stop_head` is built iff `build_stop_head` —
        only `Stage2Autoregressive` ever passes `True` (`n_sec.mode ==
        "stop_token"`, mutually exclusive with `n_sec_head`), a single
        `cond_out_dim -> 1` logit per call, same `HeadConfig` shape rules as
        the other two heads.
        """
        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)
        self.stop_head = None
        if build_stop_head:
            head_cfg = HeadConfig.from_dict(stop_head_cfg)
            hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
            self.stop_head = build_mlp_head(cond_out_dim, 1, hidden, head_cfg.depth)

    def _build_context_fusion(self, x_dim: int, context_dim: int, cond_out_dim: int) -> None:
        """Builds `self.context_adapter`/`self.fuse` — the stage-2-style
        context-fusion pattern (project the previous stage's outcome down to
        `context_dim` via `ContextAdapter`, concat onto the base conditioning,
        project back to `cond_out_dim`) shared by `Stage2OneShot` and a
        `stage="stage2"` `CriticModel` (gitea #57). Call from a subclass's
        `__init__` before using `_cond_embed`."""
        self.context_adapter = ContextAdapter(x_dim, context_dim)
        self.fuse = nn.Sequential(
            nn.Linear(cond_out_dim + context_dim, cond_out_dim),
            nn.SiLU(),
        )

    def _cond_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
        """Fuses base conditioning with the previous stage's outcome — pairs
        with `_build_context_fusion`."""
        base = self.cond_enc(cond_cont, cond_cat)
        ctx = self.context_adapter(stage1_out)
        return self.fuse(torch.cat([base, ctx], dim=-1))

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

    def _require_stop_head(self) -> None:
        if self.stop_head is None:
            raise RuntimeError(
                f"this {type(self).__name__} has no stop_head — only a "
                "Stage2Autoregressive built with stage2_model.n_sec.mode = "
                "'stop_token' owns one"
            )


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,
        mat_vocab: int,
        particle_cfg: ConditioningAxisConfig,
        material_cfg: ConditioningAxisConfig,
        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,
        trunk_type: str = "resmlp",
        block_conditioning: str = "add",
        n_sec_head_k_max: int | None = None,
        cond_enc: ConditionEncoder | None = None,
        n_sec_head_cfg: dict | None = None,
    ) -> None:
        super().__init__(
            pdg_vocab,
            mat_vocab,
            particle_cfg,
            material_cfg,
            cond_out_dim=cond_out_dim,
            generator=generator,
            noise_dim=noise_dim,
            cond_enc=cond_enc,
        )
        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,
        )

    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 _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(StageModel):
    """Predicts all `k_max` secondary slots simultaneously — v0.2 behaviour,
    reproduced exactly (`decoder = "autoregressive"` is `Stage2Autoregressive`,
    step 4/5, not implemented yet).

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

    `particle_type_cfg.target` (default `"physical"`) selects the
    secondary-type mechanism: `"physical"` keeps the type slice folded into
    the trunk's own
    flow-matched/WGAN output, unchanged from v0.2 (`sec_dim` — computed by
    the caller via `stage2_trunk_sec_dim` — already reflects this). Under
    `"onehot"`/`"embedding"` with an objective (`giant.model.objectives`) that
    doesn't fold the type slice (flow/ddpm), the type
    slice is predicted by a separate `type_head` instead (same shape pattern
    as `n_sec_head`) — `sec_dim` then covers only the continuous
    stick/dir slots, `type_head` covers `k_max * emb_dim` type logits/vectors.
    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.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        particle_cfg: ConditioningAxisConfig,
        material_cfg: ConditioningAxisConfig,
        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,
        trunk_type: str = "resmlp",
        block_conditioning: str = "add",
        build_n_sec_head: bool = True,
        particle_type_cfg: ParticleTypeConfig | None = None,
        cond_enc: ConditionEncoder | None = None,
        n_sec_head_cfg: dict | None = None,
        type_head_cfg: dict | None = None,
    ) -> None:
        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._build_context_fusion(x_dim, context_dim, cond_out_dim)
        target = self.particle_type_cfg.target
        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 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:
        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)

    def predict_type(
        self,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor,
    ) -> torch.Tensor:
        """`(B, k_max, emb_dim)` per-slot type logits (`target="onehot"`) or
        vectors (`target="embedding"`) — only under `generator in ("flow",
        "ddpm")`; `generator == "wgan"` folds the type slice into `forward`'s
        own output instead (see class docstring)."""
        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.k_max, self.type_dim)


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
    `AttentionHistory` (`attn_n_heads`/`attn_n_layers`, attention only).
    `teacher_forcing` handling lives entirely in the trainer
    (`giant/train.py`), since it only affects how training inputs are
    assembled, not this module's architecture.

    Under teacher forcing every token's conditioning is built from ground
    truth, so a whole K-token sequence trains in one parallel batched pass:
    `forward` accepts `(B, K, ...)` tensors for an arbitrary K (not hardcoded
    to `k_max`) — this also means a future one-token-at-a-time inference loop
    (`K=1` per call, step 6) needs no interface change here.

    Two independent conditioning paths, mirroring `Stage2OneShot`'s
    `_cond_embed` but split in two: `_base_cond` (`cond_enc` +
    `context_adapter` only) feeds `predict_n_sec`, since n_sec doesn't depend
    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`/`predict_stop`/
    the trunk.

    `n_sec.mode = "stop_token"` (`build_stop_head=True`) replaces
    `predict_n_sec`'s one-shot classifier with `predict_stop`'s per-token EOS
    logit instead — the two heads are mutually exclusive (`build_n_sec_head`
    is `False` whenever this is `True`, see `giant.model.builders`).
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        particle_cfg: ConditioningAxisConfig,
        material_cfg: ConditioningAxisConfig,
        hidden_dim: int = 256,
        n_res_blocks: int = 6,
        cond_out_dim: int = 128,
        context_dim: int = 64,
        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,
        trunk_type: str = "resmlp",
        block_conditioning: str = "add",
        build_n_sec_head: bool = True,
        particle_type_cfg: ParticleTypeConfig | None = None,
        history: str = "markov",
        attn_n_heads: int = 4,
        attn_n_layers: int = 2,
        cond_enc: ConditionEncoder | None = None,
        n_sec_head_cfg: dict | None = None,
        type_head_cfg: dict | None = None,
        build_stop_head: bool = False,
        stop_sampling: str = "greedy",
        stop_head_cfg: dict | None = None,
    ) -> None:
        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.stop_sampling = stop_sampling
        self.context_adapter = ContextAdapter(x_dim, context_dim)
        self.base_fuse = nn.Sequential(
            nn.Linear(cond_out_dim + context_dim, cond_out_dim),
            nn.SiLU(),
        )

        # Reuses conditioning.out_dim for the history encoder's own output
        # width — there's no dedicated stage2_model.autoregressive key for
        # this, a reasonable default rather than a design-doc-specified value.
        history_dim = cond_out_dim
        hist_in_dim = CONT_SLOT_DIM + self.type_dim
        self.history_encoder: HistoryEncoder = build_history(
            history, hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers
        )
        token_fuse_in = cond_out_dim + context_dim + history_dim + 2  # +2: remaining_frac, slot_idx
        self.token_fuse = nn.Sequential(
            nn.Linear(token_fuse_in, cond_out_dim),
            nn.SiLU(),
        )

        # `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.target
        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,
            build_stop_head=build_stop_head,
            stop_head_cfg=stop_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)
        ctx = self.context_adapter(stage1_out)
        return self.base_fuse(torch.cat([base, ctx], dim=-1))

    def _token_cond(
        self,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor,
        history_feat: torch.Tensor,
        has_prev: torch.Tensor,
        remaining_frac: torch.Tensor,
        slot_idx: torch.Tensor,
        hist: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """`hist`, if given, overrides recomputing `self.history_encoder`
        from `history_feat`/`has_prev` — the inference-time KV-cache path
        (`Stage2Autoregressive.history_step`) precomputes it once per slot and
        passes it in here so a slot's (possibly several) model calls — an ODE
        loop's substeps, or a separate `predict_type` call — read the same
        cached history instead of each re-deriving (and, under attention,
        re-appending to the cache — see `AttentionHistory.step`'s docstring)."""
        K = history_feat.size(1)
        base = self.cond_enc(cond_cont, cond_cat).unsqueeze(1).expand(-1, K, -1)
        ctx = self.context_adapter(stage1_out).unsqueeze(1).expand(-1, K, -1)
        if hist is None:
            hist = self.history_encoder(history_feat, has_prev)
        scalars = torch.stack([remaining_frac, slot_idx], dim=-1)
        return self.token_fuse(torch.cat([base, ctx, hist, scalars], dim=-1))

    def init_history_cache(self):
        """Inference-only incremental-decoding state for `self.history_encoder`
        (`giant/sample.py`'s AR loop) — whatever `self.history_encoder.init_cache()`
        returns for the configured `history` type: `None` under `history="markov"`
        (its per-step cost is already O(1) — see `HistoryEncoder`'s docstring),
        or `AttentionHistory.init_cache()`'s real per-block KV cache under
        `history="attention"`."""
        return self.history_encoder.init_cache()

    def history_step(self, token_feat: torch.Tensor, has_prev: torch.Tensor, cache) -> tuple[torch.Tensor, object]:
        """One inference slot's worth of history encoding: advances `cache`
        (from `init_history_cache`, or a previous `history_step` call) by
        `token_feat`/`has_prev` (`(B, 1, ...)` — the just-emitted previous
        token, same convention `giant.sample.sample_secondaries_ar` already
        threads as `prev_repr`), and returns `(hist, new_cache)` — `hist` is
        this slot's history summary (pass it as `_token_cond`'s `hist=` to
        every model call made for this slot), `new_cache` is what to pass into
        the *next* slot's `history_step`. Must be called exactly once per
        slot — see `AttentionHistory.step`'s docstring."""
        return self.history_encoder.step(token_feat, has_prev, cache)

    def forward(
        self,
        x_t: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor,
        history_feat: torch.Tensor,
        has_prev: torch.Tensor,
        remaining_frac: torch.Tensor,
        slot_idx: torch.Tensor,
        t: torch.Tensor | None = None,
        hist: torch.Tensor | None = None,
    ) -> torch.Tensor:
        B, K = x_t.shape[0], x_t.shape[1]
        c_emb = self._token_cond(
            cond_cont,
            cond_cat,
            stage1_out,
            history_feat,
            has_prev,
            remaining_frac,
            slot_idx,
            hist=hist,
        )
        if self.time_emb is not None:
            assert t is not None
            t_emb = self.time_emb(t.reshape(-1)).view(B, K, -1)
            cond = torch.cat([t_emb, c_emb], dim=-1)
        else:
            cond = c_emb
        x_flat = x_t.reshape(B * K, -1)
        cond_flat = cond.reshape(B * K, -1)
        cond_cont_flat = cond_cont.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1)
        cond_cat_flat = cond_cat.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1)
        out = self.trunk(x_flat, cond_flat, cond_cont_flat, cond_cat_flat)
        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:
        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(
        self,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor,
        history_feat: torch.Tensor,
        has_prev: torch.Tensor,
        remaining_frac: torch.Tensor,
        slot_idx: torch.Tensor,
        hist: torch.Tensor | None = None,
    ) -> torch.Tensor:
        self._require_type_head()
        assert self.type_head is not None
        c_emb = self._token_cond(
            cond_cont,
            cond_cat,
            stage1_out,
            history_feat,
            has_prev,
            remaining_frac,
            slot_idx,
            hist=hist,
        )
        B, K, _ = c_emb.shape
        return self.type_head(c_emb.reshape(B * K, -1)).view(B, K, self.type_dim)

    def predict_stop(
        self,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor,
        history_feat: torch.Tensor,
        has_prev: torch.Tensor,
        remaining_frac: torch.Tensor,
        slot_idx: torch.Tensor,
        hist: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """`(B, K)` raw stop logits — `n_sec.mode = "stop_token"` only.
        Evaluated on slot `k`'s own conditioning (which carries slot `k-1`'s
        history, same as `predict_type`), so this is `P(n_sec == k |
        prefix)`: a high logit at slot `k` means "stop before generating a
        token here" — the caller (`giant.sample.sample_secondaries_ar`)
        checks it before spending a model call on that slot's token."""
        self._require_stop_head()
        assert self.stop_head is not None
        c_emb = self._token_cond(
            cond_cont,
            cond_cat,
            stage1_out,
            history_feat,
            has_prev,
            remaining_frac,
            slot_idx,
            hist=hist,
        )
        B, K, _ = c_emb.shape
        return self.stop_head(c_emb.reshape(B * K, -1)).view(B, K)


class CriticModel(StageModel):
    """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`, via `StageModel._build_context_fusion`/`_cond_embed`).
    Used only when that stage's `generator == "wgan"`.

    Subclasses `StageModel` for the `cond_enc` construction and (stage 2)
    context-fusion scaffolding only — its trunk is built directly via
    `build_trunk` (output width 1) rather than through
    `_build_trunk_and_heads`, since that helper is shaped around a
    generator's `Objective`/time-embedding/flow-matching concerns
    (`forward`'s `(x_t, cond) -> vector` shape) that don't apply to a critic's
    `(x, cond) -> scalar` (gitea #57). `generator="wgan"` is passed to the
    base purely because that's factually when a critic exists; nothing here
    ever calls `_build_trunk_and_heads`, so no head/time-embedding machinery
    is built from it. Never routed (MoE) — that's a separate, unrequested
    axis of scope; see gitea #57's proposal, which covers only the trunk/
    block registries."""

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        particle_cfg: ConditioningAxisConfig,
        material_cfg: ConditioningAxisConfig,
        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,
        trunk_type: str = "resmlp",
        block_conditioning: str = "add",
    ) -> None:
        super().__init__(
            pdg_vocab,
            mat_vocab,
            particle_cfg,
            material_cfg,
            cond_out_dim=cond_out_dim,
            generator="wgan",
            noise_dim=0,
        )
        if stage not in ("stage1", "stage2"):
            raise ValueError(f"stage must be 'stage1' or 'stage2', got {stage!r}")
        self.stage = stage
        if stage == "stage2":
            self._build_context_fusion(context_in_dim, context_dim, cond_out_dim)
        self.trunk = build_trunk(
            None, trunk_type, in_dim, 1, hidden_dim, n_res_blocks, cond_out_dim, dropout, block_conditioning
        )

    def forward(
        self,
        x: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if self.stage == "stage2":
            assert stage1_out is not None, "stage='stage2' CriticModel requires stage1_out"
            cond = self._cond_embed(cond_cont, cond_cat, stage1_out)
        else:
            cond = self.cond_enc(cond_cont, cond_cat)
        return self.trunk(x, cond, cond_cont, cond_cat).squeeze(-1)