giant/giant/model/layers.py

"""Small stateless-ish building blocks shared across encoders/trunks/models —
no dependency on any other `giant.model` submodule (issues.md Issue 8)."""

import math

import torch
import torch.nn as nn


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)


def build_mlp_head(
    in_dim: int, out_dim: int, hidden: int, depth: int = 2, act: type[nn.Module] = nn.SiLU
) -> nn.Sequential:
    """`depth`-layer MLP head (gitea #36) — factors out the n_sec_head/
    type_head pattern duplicated five times across `giant.model.models`.

    `depth=1` is a bare `Linear(in_dim, out_dim)` (no hidden layer/
    activation); `depth>=2` is `Linear(in_dim, hidden) -> act -> [Linear
    (hidden, hidden) -> act] * (depth-2) -> Linear(hidden, out_dim)` —
    `depth=2` reproduces every pre-#36 n_sec_head/type_head exactly when
    `hidden == hidden_dim // 2`. Mirrors `_make_axis_mlp`'s depth
    convention above, but takes `hidden` and `out_dim` as independent
    widths (n_sec_head/type_head's hidden width is not their output width,
    unlike the particle/material axis MLPs)."""
    if depth < 1:
        raise ValueError(f"depth must be >= 1, got {depth}")
    if depth == 1:
        return nn.Sequential(nn.Linear(in_dim, out_dim))
    layers: list[nn.Module] = [nn.Linear(in_dim, hidden), act()]
    for _ in range(depth - 2):
        layers += [nn.Linear(hidden, hidden), act()]
    layers.append(nn.Linear(hidden, out_dim))
    return nn.Sequential(*layers)


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

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


BLOCK_REGISTRY: dict[str, type[nn.Module]] = {}


def register_block(name: str):
    def decorator(cls: type[nn.Module]) -> type[nn.Module]:
        BLOCK_REGISTRY[name] = cls
        return cls

    return decorator


def build_block(name: str, dim: int, cond_dim: int, dropout: float = 0.0) -> nn.Module:
    """Factory: look up a registered conditioning-injection block by name and
    construct one instance — `trunk.block_conditioning` (gitea #34)."""
    if name not in BLOCK_REGISTRY:
        raise ValueError(f"unknown block conditioning type {name!r}; available: {sorted(BLOCK_REGISTRY)}")
    return BLOCK_REGISTRY[name](dim, cond_dim, dropout)


@register_block("add")
class ResBlock(nn.Module):
    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        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


@register_block("film")
class FilmResBlock(nn.Module):
    """FiLM conditioning (Perez et al. 2018): a per-channel scale+shift
    modulates the normalized features, on top of the norm's own affine —
    an *additional* modulation, unlike `AdaLNResBlock` below, which replaces
    the norm's affine outright. `film_proj` is zero-initialized so
    `gamma=beta=0` at construction — conditioning has no effect on the
    output until training moves it, a stable starting point (though not a
    literal identity block, since `linear1`/`linear2` aren't zero-init)."""

    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.linear1 = nn.Linear(dim, dim)
        self.film_proj = nn.Linear(cond_dim, 2 * dim)
        nn.init.zeros_(self.film_proj.weight)
        nn.init.zeros_(self.film_proj.bias)
        self.act = nn.SiLU()
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim, dim)

    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        h = self.norm(x)
        gamma, beta = self.film_proj(cond).chunk(2, dim=-1)
        h = h * (1 + gamma) + beta
        h = self.linear1(h)
        h = self.act(h)
        h = self.dropout(h)
        h = self.linear2(h)
        return x + h


@register_block("adaln")
class AdaLNResBlock(nn.Module):
    """AdaLN-Zero conditioning (DiT, Peebles & Xie 2022): the norm's own
    affine is replaced by a conditioning-derived scale/shift, and the
    residual branch is scaled by a conditioning-derived gate. `adaln_proj`
    is zero-initialized, so `scale=shift=gate=0` at construction — the block
    is the exact identity function at init (`x + 0 * h' == x`), regardless
    of `x`/`cond`."""

    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(dim, elementwise_affine=False)
        self.linear1 = nn.Linear(dim, dim)
        self.adaln_proj = nn.Linear(cond_dim, 3 * dim)
        nn.init.zeros_(self.adaln_proj.weight)
        nn.init.zeros_(self.adaln_proj.bias)
        self.act = nn.SiLU()
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim, dim)

    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        h = self.norm(x)
        scale, shift, gate = self.adaln_proj(cond).chunk(3, dim=-1)
        h = h * (1 + scale) + shift
        h = self.linear1(h)
        h = self.act(h)
        h = self.dropout(h)
        h = self.linear2(h)
        return x + gate * h