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.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 AttentionHistory, HistoryEncoder, MarkovHistory
from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding
from giant.model.routers import Router
from giant.model.trunks import build_trunk

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


def resolve_type_n_classes(particle_type_cfg: dict, 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.get("target", "physical") == "onehot":
        return particle_type_cfg.get("n_classes", 0) or particle_emb_dim
    return particle_emb_dim


def stage2_type_dim(particle_type_cfg: dict, 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)."""
    target = particle_type_cfg.get("target", "physical")
    return PARTICLE_PHYS_DIM if target == "physical" else emb_dim


def stage2_trunk_sec_dim(particle_type_cfg: dict, 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 `generator == "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)`. Under
    `generator in ("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`.
    """
    target = particle_type_cfg.get("target", "physical")
    if target == "physical":
        return k_max * SEC_SLOT_DIM
    if generator == "wgan":
        return k_max * (CONT_SLOT_DIM + emb_dim)
    return k_max * CONT_SLOT_DIM


class Stage1Model(nn.Module):
    """Predicts the 9D primary post-step vector. No `n_sec_head` — fresh runs
    move it to stage 2, except for a migrated v0.2 checkpoint
    (`n_sec_head_k_max` given), where it stays attached here
    since that's where its weights live and what conditioning it was trained
    against (see `_migrate_legacy_model_config`).

    `cond_enc`, if given, is used in place of building a fresh
    `ConditionEncoder` — `conditioning.share_stages = true`: `build_models`
    constructs one shared instance and passes it to both stages, halving the
    conditioning parameter count and forcing a common representation."""

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        particle_cfg: dict,
        material_cfg: dict,
        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",
        n_sec_head_k_max: int | None = None,
        cond_enc: ConditionEncoder | None = None,
    ) -> None:
        super().__init__()
        self.generator_kind = generator
        self.noise_dim = noise_dim
        self.cond_enc = (
            cond_enc
            if cond_enc is not None
            else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
        )
        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, trunk_type, 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."""
        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')"
            )
        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 (`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 `generator in ("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 `generator == "wgan"` the type slice stays folded into `sec_dim`
    (just `emb_dim` instead of `PARTICLE_PHYS_DIM` wide) and `type_head` is
    unused (`None`) — the WGAN trainer handles the ST-Gumbel relaxation.

    `cond_enc`, if given, is used in place of building a fresh
    `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`).
    """

    def __init__(
        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,
        trunk_type: str = "resmlp",
        build_n_sec_head: bool = True,
        particle_type_cfg: dict | None = None,
        cond_enc: ConditionEncoder | None = None,
    ) -> None:
        super().__init__()
        self.generator_kind = generator
        self.noise_dim = noise_dim
        self.k_max = k_max
        self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"})
        self.type_dim = stage2_type_dim(
            self.particle_type_cfg, resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"])
        )
        self.cond_enc = (
            cond_enc
            if cond_enc is not None
            else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
        )
        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, trunk_type, 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),
            )
        self.type_head = None
        target = self.particle_type_cfg.get("target", "physical")
        if target != "physical" and generator in ("flow", "ddpm"):
            emb_dim = resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"])
            self.type_head = nn.Sequential(
                nn.Linear(cond_out_dim, hidden_dim // 2),
                nn.SiLU(),
                nn.Linear(hidden_dim // 2, k_max * emb_dim),
            )
            self._type_k_max = k_max
            self._type_emb_dim = emb_dim

    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 (n_sec.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)

    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)."""
        if self.type_head is None:
            raise RuntimeError(
                "this Stage2OneShot has no type_head — either "
                "particle_type.target='physical' (the type slice is part of "
                "forward()'s own output) or generator='wgan' (the WGAN "
                "trainer reads the type slice out of forward()'s output "
                "directly instead)"
            )
        c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
        return self.type_head(c_emb).view(-1, self._type_k_max, self._type_emb_dim)


class Stage2Autoregressive(nn.Module):
    """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`/the trunk.

    `cond_enc`, if given, is used in place of building a fresh
    `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`).
    """

    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,
        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",
        build_n_sec_head: bool = True,
        particle_type_cfg: dict | None = None,
        history: str = "markov",
        attn_n_heads: int = 4,
        attn_n_layers: int = 2,
        cond_enc: ConditionEncoder | None = None,
    ) -> None:
        super().__init__()
        if history not in ("markov", "attention"):
            raise ValueError(f"stage2_model.autoregressive.history={history!r} — must be 'markov' or 'attention'")
        self.history_kind = history
        self.generator_kind = generator
        self.noise_dim = noise_dim
        self.k_max = k_max
        self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"})
        emb_dim = resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"])
        self.type_dim = stage2_type_dim(self.particle_type_cfg, emb_dim)

        self.cond_enc = (
            cond_enc
            if cond_enc is not None
            else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
        )
        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 = (
            AttentionHistory(hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers)
            if history == "attention"
            else MarkovHistory(hist_in_dim, history_dim)
        )
        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(),
        )

        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
        token_dim = stage2_trunk_sec_dim(self.particle_type_cfg, generator, 1, emb_dim)
        in_dim = noise_dim if generator == "wgan" else token_dim
        self.trunk = build_trunk(
            router,
            trunk_type,
            in_dim,
            token_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),
            )
        self.type_head = None
        target = self.particle_type_cfg.get("target", "physical")
        if target != "physical" and generator in ("flow", "ddpm"):
            self.type_head = nn.Sequential(
                nn.Linear(cond_out_dim, hidden_dim // 2),
                nn.SiLU(),
                nn.Linear(hidden_dim // 2, self.type_dim),
            )

    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): `None` under `history="markov"` (its
        per-step cost is already O(1) — see `HistoryEncoder`'s docstring), or
        `AttentionHistory.init_cache()` under `history="attention"`."""
        if isinstance(self.history_encoder, AttentionHistory):
            return self.history_encoder.init_cache()
        return None

    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."""
        if isinstance(self.history_encoder, AttentionHistory):
            return self.history_encoder.step(token_feat, has_prev, cache)
        return self.history_encoder(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:
        if self.n_sec_head is None:
            raise RuntimeError(
                "this Stage2Autoregressive has no n_sec_head — it belongs to "
                "a migrated v0.2 checkpoint (n_sec.owner='stage1'); call "
                "stage1.predict_n_sec(cond_cont, cond_cat) instead"
            )
        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:
        if self.type_head is None:
            raise RuntimeError(
                "this Stage2Autoregressive has no type_head — either "
                "particle_type.target='physical' (the type slice is part of "
                "forward()'s own output) or generator='wgan' (the WGAN "
                "trainer reads the type slice out of forward()'s output "
                "directly instead)"
            )
        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)


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)