code/rtm_core.py

import math
from typing import List, Tuple
from dataclasses import dataclass

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


# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------

def _trunc_normal_init(tensor: torch.Tensor, std: float = 1.0,
                       lower: float = -2.0, upper: float = 2.0) -> torch.Tensor:
    with torch.no_grad():
        if std == 0:
            tensor.zero_()
        else:
            sqrt2 = math.sqrt(2)
            a = math.erf(lower / sqrt2)
            b = math.erf(upper / sqrt2)
            z = (b - a) / 2
            c = (2 * math.pi) ** -0.5
            pdf_u = c * math.exp(-0.5 * lower ** 2)
            pdf_l = c * math.exp(-0.5 * upper ** 2)
            comp_std = std / math.sqrt(
                1 - (upper * pdf_u - lower * pdf_l) / z
                - ((pdf_u - pdf_l) / z) ** 2
            )
            tensor.uniform_(a, b)
            tensor.erfinv_()
            tensor.mul_(sqrt2 * comp_std)
            tensor.clip_(lower * comp_std, upper * comp_std)
    return tensor


def _find_multiple(a: int, b: int) -> int:
    return (-(a // -b)) * b


def rms_norm(hidden_states: torch.Tensor, variance_epsilon: float) -> torch.Tensor:
    input_dtype = hidden_states.dtype
    hidden_states = hidden_states.to(torch.float32)
    variance = hidden_states.square().mean(-1, keepdim=True)
    hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon)
    return hidden_states.to(input_dtype)


class CastedLinear(nn.Module):
    def __init__(self, in_features: int, out_features: int, bias: bool):
        super().__init__()
        self.weight = nn.Parameter(
            _trunc_normal_init(torch.empty(out_features, in_features),
                               std=1.0 / (in_features ** 0.5))
        )
        self.bias = nn.Parameter(torch.zeros(out_features)) if bias else None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.linear(
            x, self.weight.to(x.dtype),
            bias=self.bias.to(x.dtype) if self.bias is not None else None,
        )


class SwiGLU(nn.Module):
    def __init__(self, hidden_size: int, expansion: float):
        super().__init__()
        inter = _find_multiple(round(expansion * hidden_size * 2 / 3), 256)
        self.gate_up_proj = CastedLinear(hidden_size, inter * 2, bias=False)
        self.down_proj = CastedLinear(inter, hidden_size, bias=False)

    def forward(self, x):
        gate, up = self.gate_up_proj(x).chunk(2, dim=-1)
        return self.down_proj(F.silu(gate) * up)


# ---------------------------------------------------------------------------
# RTM block: token-mixing MLP + channel-mixing MLP with RMSNorm residuals
# ---------------------------------------------------------------------------

class RTMBlock(nn.Module):

    def __init__(self, hidden_size: int, seq_len: int, expansion: float,
                 rms_norm_eps: float = 1e-5):
        super().__init__()
        self.mlp_t = SwiGLU(hidden_size=seq_len, expansion=expansion)
        self.mlp = SwiGLU(hidden_size=hidden_size, expansion=expansion)
        self.norm_eps = rms_norm_eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = hidden_states.transpose(1, 2)
        hidden_states = rms_norm(
            hidden_states + self.mlp_t(hidden_states),
            variance_epsilon=self.norm_eps,
        )
        hidden_states = hidden_states.transpose(1, 2)
        hidden_states = rms_norm(
            hidden_states + self.mlp(hidden_states),
            variance_epsilon=self.norm_eps,
        )
        return hidden_states


class RTMReasoningModule(nn.Module):

    def __init__(self, layers: List[RTMBlock]):
        super().__init__()
        self.layers = nn.ModuleList(layers)

    def forward(self, hidden_states: torch.Tensor,
                input_injection: torch.Tensor) -> torch.Tensor:
        hidden_states = hidden_states + input_injection
        for layer in self.layers:
            hidden_states = layer(hidden_states)
        return hidden_states


def _make_level(num_layers: int, hidden_size: int, seq_len: int,
                expansion: float, rms_norm_eps: float) -> RTMReasoningModule:
    return RTMReasoningModule([
        RTMBlock(hidden_size, seq_len, expansion, rms_norm_eps)
        for _ in range(num_layers)
    ])


# ---------------------------------------------------------------------------
# Recursive H/L cycle loop
# ---------------------------------------------------------------------------

@dataclass
class RTMCarry:
    z_H: torch.Tensor
    z_L: torch.Tensor


class RTMInner(nn.Module):

    def __init__(self, hidden_size: int, expansion: float,
                 H_cycles: int, L_cycles: int, H_layers: int, L_layers: int,
                 num_tokens: int,
                 with_grad: bool = False,
                 cycle_noise_std: float = 0.0,
                 rms_norm_eps: float = 1e-5,
                 forward_dtype: str = "float32"):
        super().__init__()

        self.hidden_size = hidden_size
        self.H_cycles = H_cycles
        self.L_cycles = L_cycles
        self.with_grad = with_grad
        self.cycle_noise_std = max(0.0, float(cycle_noise_std))
        self.forward_dtype = getattr(torch, forward_dtype)

        self.total_seq_len = max(1, int(num_tokens))

        self.L_level = _make_level(L_layers, hidden_size, self.total_seq_len,
                                   expansion, rms_norm_eps)
        self.H_init = nn.Buffer(
            _trunc_normal_init(torch.empty(hidden_size, dtype=self.forward_dtype), std=1),
            persistent=True,
        )
        self.L_init = nn.Parameter(
            _trunc_normal_init(torch.empty(hidden_size, dtype=self.forward_dtype), std=1)
        )

    def empty_carry(self, batch_size: int, device=None) -> RTMCarry:
        if device is None:
            device = self.H_init.device
        return RTMCarry(
            z_H=self.H_init.unsqueeze(0).unsqueeze(0).expand(
                batch_size, self.total_seq_len, -1),
            z_L=self.L_init.unsqueeze(0).unsqueeze(0).expand(
                batch_size, self.total_seq_len, -1),
        )

    def forward(self, carry: RTMCarry, z_H_init: torch.Tensor
                ) -> Tuple[RTMCarry, torch.Tensor, List[torch.Tensor]]:
        z_H, z_L = carry.z_H, carry.z_L
        intermediates: List[torch.Tensor] = []

        if self.with_grad:
            for _ in range(self.H_cycles):
                for _ in range(self.L_cycles):
                    z_L = self.L_level(z_L, z_H + z_H_init)
                z_H = self.L_level(z_H, z_L)
                if self.training and self.cycle_noise_std > 0:
                    z_H = z_H + torch.randn_like(z_H) * self.cycle_noise_std
                intermediates.append(z_H)
        else:
            with torch.no_grad():
                for _ in range(self.H_cycles - 1):
                    for _ in range(self.L_cycles):
                        z_L = self.L_level(z_L, z_H + z_H_init)
                    z_H = self.L_level(z_H, z_L)
                    if self.training and self.cycle_noise_std > 0:
                        z_H = z_H + torch.randn_like(z_H) * self.cycle_noise_std
                    intermediates.append(z_H)

            for _ in range(self.L_cycles):
                z_L = self.L_level(z_L, z_H + z_H_init)
            z_H = self.L_level(z_H, z_L)
            if self.training and self.cycle_noise_std > 0:
                z_H = z_H + torch.randn_like(z_H) * self.cycle_noise_std
            intermediates.append(z_H)

        new_carry = RTMCarry(z_H=z_H.detach(), z_L=z_L.detach())
        return new_carry, z_H, intermediates

class _PixelNorm(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * (x.square().mean(dim=1, keepdim=True) + 1e-8).rsqrt()


class _EqualLinear(nn.Module):
    def __init__(self, in_dim: int, out_dim: int):
        super().__init__()
        self.linear = nn.Linear(in_dim, out_dim)
        self.linear.bias.data.zero_()

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


def _coerce_noise_2d(input_tensor: torch.Tensor, code_dim: int, mapper_name: str) -> torch.Tensor:
    """Accept [B, D] or [B, 1, D] and return [B, D]."""
    if input_tensor.dim() == 2 and input_tensor.shape[1] == code_dim:
        return input_tensor
    if input_tensor.dim() == 3 and input_tensor.shape[1] == 1 and input_tensor.shape[2] == code_dim:
        return input_tensor.squeeze(1)
    raise ValueError(
        f"{mapper_name} expected input shape [B, {code_dim}] or [B, 1, {code_dim}], "
        f"got {tuple(input_tensor.shape)}"
    )


class RTMMappingNetwork(nn.Module):
    def __init__(self, code_dim: int,
                 num_tokens: int = 1,
                 H_cycles: int = 1, L_cycles: int = 1,
                 H_layers: int = 2, L_layers: int = 2,
                 hidden_size: int = 256,
                 expansion: float = 4.0, refinement_steps: int = 1,
                 with_grad: bool = False,
                 cycle_noise_std: float = 0.0,
                 rms_norm_eps: float = 1e-5,
                 forward_dtype: str = "float32"):
        super().__init__()
        self.code_dim = code_dim
        self.refinement_steps = max(1, refinement_steps)

        total_seq_len = max(1, int(num_tokens))

        if hidden_size <= 0:
            assert code_dim % total_seq_len == 0, (
                f"code_dim={code_dim} must be divisible by num_tokens={total_seq_len} "
                f"when hidden_size is auto"
            )
            hidden_size = code_dim // total_seq_len

        self.hidden_size = hidden_size

        self.trm = RTMInner(
            hidden_size=hidden_size,
            expansion=expansion,
            H_cycles=H_cycles,
            L_cycles=L_cycles,
            H_layers=H_layers,
            L_layers=L_layers,
            num_tokens=num_tokens,
            with_grad=with_grad,
            cycle_noise_std=cycle_noise_std,
            rms_norm_eps=rms_norm_eps,
            forward_dtype=forward_dtype,
        )

        seq_len = self.trm.total_seq_len
        capacity = seq_len * hidden_size

        self.pixel_norm = _PixelNorm()
        self.mapper_direct = (capacity == code_dim)
        if not self.mapper_direct:
            self.z_to_seq = _EqualLinear(code_dim, capacity)
            self.seq_to_w = _EqualLinear(capacity, code_dim)

    def _z_H_to_w(self, z_H: torch.Tensor) -> torch.Tensor:
        flat = z_H.flatten(start_dim=1)
        if self.mapper_direct:
            return flat
        return self.seq_to_w(flat)

    def forward(self, input):
        if isinstance(input, (list, tuple)):
            input = input[0]

        input = _coerce_noise_2d(input, self.code_dim, self.__class__.__name__)
        B = input.shape[0]
        seq_len = self.trm.total_seq_len

        z_norm = self.pixel_norm(input)
        if self.mapper_direct:
            z_seq = z_norm.view(B, seq_len, self.hidden_size)
        else:
            z_seq = self.z_to_seq(z_norm).view(B, seq_len, self.hidden_size)

        carry = self.trm.empty_carry(B, device=input.device)
        for _ in range(self.refinement_steps):
            carry, z_H_out, _ = self.trm(carry, z_H_init=z_seq)

        return [self._z_H_to_w(z_H_out)]

    def forward_w_trajectory(self, input):
        """Decode every intermediate H state into a style vector.

        Returns ``[w_proj, w_after_cycle_1, ...]`` so callers can visualize
        how w is refined across the H/L cycles.
        """
        if isinstance(input, (list, tuple)):
            input = input[0]

        input = _coerce_noise_2d(input, self.code_dim, self.__class__.__name__)
        B = input.shape[0]
        seq_len = self.trm.total_seq_len

        z_norm = self.pixel_norm(input)
        if self.mapper_direct:
            z_seq = z_norm.view(B, seq_len, self.hidden_size)
        else:
            z_seq = self.z_to_seq(z_norm).view(B, seq_len, self.hidden_size)

        trajectory = [self._z_H_to_w(z_seq)]
        carry = self.trm.empty_carry(B, device=input.device)
        for _ in range(self.refinement_steps):
            carry, _z_H_out, intermediates = self.trm(carry, z_H_init=z_seq)
            for zh in intermediates:
                trajectory.append(self._z_H_to_w(zh))
        return trajectory