DeepHealth/backbones.py

import math

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

from attention_types import (
    DIFF_ATTENTION,
    STANDARD_ATTENTION,
    resolve_attention_type,
)
from model_architectures import (
    TRAJ_MIXER_ARCHITECTURE,
    TRANSFORMER_FFN_ARCHITECTURE,
    resolve_model_architecture,
)


class TimeRoPE(nn.Module):
    def __init__(self, dim: int, base: float = 10000.0):
        super().__init__()
        assert dim % 2 == 0, "RoPE dim must be even"
        self.dim = dim
        # inv_freq is not trainable, but should move with device.
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def precompute_cache(self, tau: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        t = tau.unsqueeze(-1)  # (B, L, 1)
        angles = t * self.inv_freq  # (B, L, dim//2)
        # Pre-expand for heads and interleave once (avoids N_layers repeats)
        cos = angles.cos().unsqueeze(1).repeat_interleave(2, dim=-1)
        sin = angles.sin().unsqueeze(1).repeat_interleave(2, dim=-1)
        return cos, sin  # (B, 1, L, dim)

    @staticmethod
    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
        """Rotate pairs: ``[-x2, x1, -x4, x3, ...]``."""
        x1 = x[..., 0::2]
        x2 = x[..., 1::2]
        return torch.stack((-x2, x1), dim=-1).flatten(-2)

    @staticmethod
    def apply_from_cache(
        q: torch.Tensor,
        k: torch.Tensor,
        rope_cache: tuple[torch.Tensor, torch.Tensor],
    ) -> tuple[torch.Tensor, torch.Tensor]:
        cos, sin = rope_cache  # each (B, 1, L, dim)
        q_rot = q * cos + TimeRoPE._rotate_half(q) * sin
        k_rot = k * cos + TimeRoPE._rotate_half(k) * sin
        return q_rot, k_rot

    def forward(
        self,
        tau: torch.Tensor,
        q: torch.Tensor,
        k: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        cache = self.precompute_cache(tau)
        return self.apply_from_cache(q, k, cache)


class GaussianRBFTimeBasis(nn.Module):
    def __init__(
        self,
        n_bases: int = 16,
        max_time_diff: float = 40.0,
    ):
        super().__init__()
        self.n_bases = n_bases

        # Evenly spaced RBF centres for non-negative linear time differences.
        # Causal masking enforces query_time >= key_time, so diff is >= 0.
        centers = torch.linspace(0.0, max_time_diff, n_bases)
        self.register_buffer("centers", centers,
                             persistent=False)  # (n_bases,)

        # Learnable log-widths (initialized to center spacing on linear scale).
        init_width = max(max_time_diff / max(n_bases - 1, 1), 1e-3)
        init_log_width = math.log(init_width)
        self.log_widths = nn.Parameter(torch.full((n_bases,), init_log_width))

    def precompute_cache(self, tau: torch.Tensor) -> torch.Tensor:

        time_coord = tau.float()  # (B, L)
        # Pairwise signed difference: query_i - key_j.
        diff = time_coord.unsqueeze(
            2) - time_coord.unsqueeze(1)  # (B, L_q, L_k)
        # Gaussian RBF: exp(-0.5 * ((diff - c) / w)^2)
        diff = diff.unsqueeze(-1)  # (B, L, L, 1)
        widths = self.log_widths.exp()  # (n_bases,)
        rbf_acts = torch.exp(
            -0.5 * ((diff - self.centers) / widths).square()
            # (B, L, L, n_bases)
        )
        return rbf_acts


class TemporalAttention(nn.Module):
    def __init__(
        self,
        n_embd: int,
        n_head: int,
        n_rbf_bases: int = 16,
        dropout: float = 0.0,
        use_time_rope: bool = True,
        use_rbf_bias: bool = True,
    ):
        super().__init__()
        assert n_embd % n_head == 0, "n_embd must be divisible by n_head"
        self.n_head = n_head
        self.d_head = n_embd // n_head
        self.scale = 1.0 / math.sqrt(self.d_head)
        self.use_time_rope = use_time_rope
        self.use_rbf_bias = use_rbf_bias

        # QKV projection (fused for efficiency)
        self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)
        # Output projection
        self.out_proj = nn.Linear(n_embd, n_embd, bias=False)

        # Layer-specific projection from shared RBF basis activations to per-head attention bias.
        self.rbf_proj = nn.Linear(n_rbf_bases, n_head, bias=False)
        # Keep the initial RBF attention bias exactly zero through the
        # zero-initialized projection, while leaving that projection with a
        # live gradient from the first optimization step.
        self.time_bias_scale = nn.Parameter(torch.tensor(1.0))

        self.resid_drop = nn.Dropout(dropout)
        self.reset_parameters()

    def reset_parameters(self) -> None:
        """Match the previous version's GPT-style weight initialization."""
        nn.init.normal_(self.qkv.weight, mean=0.0, std=0.02)
        nn.init.normal_(self.out_proj.weight, mean=0.0, std=0.02)
        nn.init.zeros_(self.rbf_proj.weight)

    def forward(
        self,
        x: torch.Tensor,
        rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
        rbf_cache: torch.Tensor | None = None,
        attn_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if self.use_time_rope:
            assert rope_cache is not None, "rope_cache must be provided when use_time_rope is True"
        if self.use_rbf_bias:
            assert rbf_cache is not None, "rbf_cache must be provided when use_rbf_bias is True"

        B, L, _ = x.shape
        H, D = self.n_head, self.d_head

        # --- QKV ----------------------------------------------------------
        qkv = self.qkv(x).reshape(B, L, 3, H, D).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)  # each (B, H, L, D)

        # --- Apply RoPE (from shared cache) --------------------------------
        if self.use_time_rope:
            q, k = TimeRoPE.apply_from_cache(q, k, rope_cache)

        # Build additive attention bias mask: time bias + causal/padding mask.
        time_bias = None
        if self.use_rbf_bias:
            time_bias = self.rbf_proj(rbf_cache).permute(
                0, 3, 1, 2)  # (B, H, L, L)
            time_bias = self.time_bias_scale.tanh() * time_bias

        if time_bias is not None and attn_mask is not None:
            attn_bias = time_bias + attn_mask.to(time_bias.dtype)
        elif time_bias is not None:
            attn_bias = time_bias
        elif attn_mask is not None:
            attn_bias = attn_mask
        else:
            attn_bias = None

        out = F.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=attn_bias,
            dropout_p=0.0,
            is_causal=False,
            scale=self.scale,
        )

        # --- Aggregate & project out --------------------------------------
        out = out.transpose(1, 2).reshape(B, L, H * D)
        return self.resid_drop(self.out_proj(out))


def diff_lambda_init(depth: int) -> float:
    """DIFF V1 layer-depth initialization (``depth`` is zero-based)."""
    if depth < 0:
        raise ValueError(f"depth must be >= 0, got {depth}")
    return 0.8 - 0.6 * math.exp(-0.3 * depth)


class DifferentialAttention(nn.Module):
    """DIFF V1 attention with paired Q/K heads and no per-head norm.

    ``n_head`` retains the baseline head count. Adjacent Q/K heads are paired,
    so the module has ``n_head // 2`` differential heads. Each paired V head
    has width ``2 * d_head`` and the merged output remains ``n_embd`` wide.
    """

    def __init__(
        self,
        n_embd: int,
        n_head: int,
        n_rbf_bases: int = 16,
        dropout: float = 0.0,
        use_time_rope: bool = True,
        use_rbf_bias: bool = True,
        depth: int = 0,
    ):
        super().__init__()
        if n_head % 2 != 0:
            raise ValueError(
                "DifferentialAttention requires an even baseline n_head, "
                f"got {n_head}"
            )
        if n_embd % n_head != 0:
            raise ValueError(
                f"n_embd must be divisible by n_head, got {n_embd} and {n_head}"
            )

        self.n_head = n_head
        self.n_diff_head = n_head // 2
        self.d_head = n_embd // n_head
        self.d_value_head = 2 * self.d_head
        self.scale = 1.0 / math.sqrt(self.d_head)
        self.use_time_rope = use_time_rope
        self.use_rbf_bias = use_rbf_bias
        self.lambda_init = diff_lambda_init(depth)

        # These projection widths intentionally match TemporalAttention.
        self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)
        self.out_proj = nn.Linear(n_embd, n_embd, bias=False)

        # Preserve one relative-time bias channel per baseline Q/K head. After
        # pairing, each differential head therefore has one bias for each map.
        self.rbf_proj = nn.Linear(n_rbf_bases, n_head, bias=False)
        self.time_bias_scale = nn.Parameter(torch.tensor(1.0))

        # DIFF V1 uses four layer-level vectors, shared across differential
        # heads. There is deliberately no per-head RMSNorm/LayerNorm here.
        self.lambda_q1 = nn.Parameter(torch.empty(self.d_head))
        self.lambda_k1 = nn.Parameter(torch.empty(self.d_head))
        self.lambda_q2 = nn.Parameter(torch.empty(self.d_head))
        self.lambda_k2 = nn.Parameter(torch.empty(self.d_head))

        self.resid_drop = nn.Dropout(dropout)
        self.reset_parameters()

    def reset_parameters(self) -> None:
        nn.init.normal_(self.qkv.weight, mean=0.0, std=0.02)
        nn.init.normal_(self.out_proj.weight, mean=0.0, std=0.02)
        nn.init.zeros_(self.rbf_proj.weight)
        nn.init.normal_(self.lambda_q1, mean=0.0, std=0.1)
        nn.init.normal_(self.lambda_k1, mean=0.0, std=0.1)
        nn.init.normal_(self.lambda_q2, mean=0.0, std=0.1)
        nn.init.normal_(self.lambda_k2, mean=0.0, std=0.1)

    def _lambda(self, dtype: torch.dtype) -> torch.Tensor:
        lambda_1 = torch.exp(
            torch.sum(self.lambda_q1 * self.lambda_k1).float()
        )
        lambda_2 = torch.exp(
            torch.sum(self.lambda_q2 * self.lambda_k2).float()
        )
        return (lambda_1 - lambda_2 + self.lambda_init).to(dtype=dtype)

    def forward(
        self,
        x: torch.Tensor,
        rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
        rbf_cache: torch.Tensor | None = None,
        attn_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if self.use_time_rope:
            assert rope_cache is not None, "rope_cache must be provided when use_time_rope is True"
        if self.use_rbf_bias:
            assert rbf_cache is not None, "rbf_cache must be provided when use_rbf_bias is True"

        batch_size, seq_len, n_embd = x.shape

        # q/k: (B, 2 * H_diff, L, d_head)
        # v:   (B, H_diff, L, 2 * d_head)
        qkv = self.qkv(x).reshape(
            batch_size, seq_len, 3, self.n_head, self.d_head
        ).permute(2, 0, 3, 1, 4)
        q, k, v_unpaired = qkv.unbind(0)
        v = v_unpaired.transpose(1, 2).reshape(
            batch_size, seq_len, self.n_diff_head, self.d_value_head
        ).transpose(1, 2)

        # A single cache is broadcast over all heads, so both Q/K maps receive
        # exactly the same TimeRoPE coordinates and rotation rule.
        if self.use_time_rope:
            q, k = TimeRoPE.apply_from_cache(q, k, rope_cache)

        time_bias = None
        if self.use_rbf_bias:
            time_bias = self.rbf_proj(rbf_cache).permute(0, 3, 1, 2)
            time_bias = self.time_bias_scale.tanh() * time_bias

        if time_bias is not None and attn_mask is not None:
            attn_bias = time_bias + attn_mask.to(time_bias.dtype)
        elif time_bias is not None:
            attn_bias = time_bias
        elif attn_mask is not None:
            attn_bias = attn_mask
        else:
            attn_bias = None

        attn_logits = torch.matmul(q, k.transpose(-1, -2)) * self.scale
        if attn_bias is not None:
            attn_logits = attn_logits + attn_bias
        attn_weights = F.softmax(
            attn_logits, dim=-1, dtype=torch.float32
        ).to(q.dtype)
        attn_weights = attn_weights.reshape(
            batch_size,
            self.n_diff_head,
            2,
            seq_len,
            seq_len,
        )
        diff_weights = (
            attn_weights[:, :, 0]
            - self._lambda(q.dtype) * attn_weights[:, :, 1]
        )

        # No V1 per-head RMSNorm and no replacement normalization: the
        # differential result goes directly to head merge and output projection.
        out = torch.matmul(diff_weights, v)
        out = out.transpose(1, 2).reshape(batch_size, seq_len, n_embd)
        return self.resid_drop(self.out_proj(out))


# Concise public alias for callers that prefer the shorter name.
DiffAttention = DifferentialAttention


def build_attention(
    attention_type: str,
    *,
    n_embd: int,
    n_head: int,
    n_rbf_bases: int = 16,
    dropout: float = 0.0,
    use_time_rope: bool = True,
    use_rbf_bias: bool = True,
    depth: int = 0,
) -> nn.Module:
    """Build a standard or DIFF V1 attention module."""
    resolved = resolve_attention_type(attention_type)
    common_kwargs = {
        "n_embd": n_embd,
        "n_head": n_head,
        "n_rbf_bases": n_rbf_bases,
        "dropout": dropout,
        "use_time_rope": use_time_rope,
        "use_rbf_bias": use_rbf_bias,
    }
    if resolved == STANDARD_ATTENTION:
        return TemporalAttention(**common_kwargs)
    if resolved == DIFF_ATTENTION:
        return DifferentialAttention(**common_kwargs, depth=depth)
    raise ValueError(f"Unsupported attention_type: {resolved!r}")


class SwiGLU(nn.Module):
    def __init__(
        self,
        n_embd: int,
        hidden_dim: int | None = None,
        dropout: float = 0.0,
        bias: bool = True,
    ):
        super().__init__()
        hidden_dim = hidden_dim if hidden_dim is not None else int(
            n_embd * 2.5)

        self.w1 = nn.Linear(n_embd, hidden_dim, bias=bias)  # gate path
        self.w2 = nn.Linear(n_embd, hidden_dim, bias=bias)  # value path
        # output projection
        self.w3 = nn.Linear(hidden_dim, n_embd, bias=bias)
        self.drop = nn.Dropout(dropout)
        self.reset_parameters()

    def reset_parameters(self) -> None:
        """GPT-style parameter initialization for MLP paths."""
        nn.init.normal_(self.w1.weight, mean=0.0, std=0.02)
        nn.init.normal_(self.w2.weight, mean=0.0, std=0.02)
        nn.init.normal_(self.w3.weight, mean=0.0, std=0.02)
        if self.w1.bias is not None:
            nn.init.zeros_(self.w1.bias)
            nn.init.zeros_(self.w2.bias)
            nn.init.zeros_(self.w3.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """``(B, L, n_embd) -> (B, L, n_embd)``."""
        return self.drop(self.w3(F.silu(self.w1(x)) * self.w2(x)))


class TrajMixer(nn.Module):
    """PreNorm gated mixing within and across latent trajectory groups.

    The groups are contiguous partitions of the post-``W_O`` residual
    representation. They are deliberately not treated as attention heads.
    All operations are position-wise, so the sequence dimension remains fully
    parallel and no temporal information can leak between positions here.
    """

    def __init__(
        self,
        n_embd: int,
        n_head: int = 10,
        dropout: float = 0.0,
    ):
        super().__init__()
        if n_embd <= 0:
            raise ValueError(f"n_embd must be > 0, got {n_embd}")
        if n_head <= 0:
            raise ValueError(f"n_head must be > 0, got {n_head}")
        if n_embd % n_head != 0:
            raise ValueError(
                f"n_embd must be divisible by n_head, got {n_embd} and {n_head}"
            )
        self.n_embd = n_embd
        self.n_group = n_head
        self.d_group = n_embd // n_head
        self.intra_hidden = 4 * self.d_group
        self.hidden_group = 4 * n_head

        self.norm = nn.LayerNorm(self.n_embd)

        self.intra_gate_proj = nn.Parameter(
            torch.empty(self.n_group, self.d_group, self.intra_hidden)
        )
        self.intra_value_proj = nn.Parameter(
            torch.empty(self.n_group, self.d_group, self.intra_hidden)
        )
        self.intra_output_proj = nn.Parameter(
            torch.empty(self.n_group, self.intra_hidden, self.d_group)
        )
        self.intra_gate_logits = nn.Parameter(
            torch.empty(self.n_group, self.d_group)
        )

        self.gate_proj = nn.Parameter(
            torch.empty(self.d_group, self.n_group, self.hidden_group)
        )
        self.value_proj = nn.Parameter(
            torch.empty(self.d_group, self.n_group, self.hidden_group)
        )
        self.output_proj = nn.Parameter(
            torch.empty(self.d_group, self.hidden_group, self.n_group)
        )
        self.drop = nn.Dropout(dropout)
        self.reset_parameters()

    def reset_parameters(self) -> None:
        for group_idx in range(self.n_group):
            nn.init.xavier_uniform_(self.intra_gate_proj[group_idx])
            nn.init.xavier_uniform_(self.intra_value_proj[group_idx])
            nn.init.xavier_uniform_(self.intra_output_proj[group_idx])
        nn.init.constant_(
            self.intra_gate_logits,
            math.log(0.1 / 0.9),
        )

        for feature_idx in range(self.d_group):
            nn.init.xavier_uniform_(self.gate_proj[feature_idx])
            nn.init.xavier_uniform_(self.value_proj[feature_idx])
        nn.init.normal_(self.output_proj, mean=0.0, std=1e-3)

    def _intra_mix(self, grouped: torch.Tensor) -> torch.Tensor:
        """Mix features independently inside each residual-space group."""
        intra_gate = torch.einsum(
            "blgd,gdh->blgh", grouped, self.intra_gate_proj
        )
        intra_value = torch.einsum(
            "blgd,gdh->blgh", grouped, self.intra_value_proj
        )
        intra_hidden = F.silu(intra_gate) * intra_value
        return torch.einsum(
            "blgh,ghd->blgd", intra_hidden, self.intra_output_proj
        )

    def _cross_mix(self, grouped: torch.Tensor) -> torch.Tensor:
        """Mix groups independently for each within-group coordinate."""
        gate = torch.einsum(
            "blgr,rgh->blhr", grouped, self.gate_proj
        )
        value = torch.einsum(
            "blgr,rgh->blhr", grouped, self.value_proj
        )
        hidden = F.silu(gate) * value
        return torch.einsum(
            "blhr,rhg->blgr", hidden, self.output_proj
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply one full-width PreNorm and one outer residual update."""
        if x.ndim != 3:
            raise ValueError(
                f"TrajMixer expects a 3D tensor, got shape {tuple(x.shape)}"
            )
        if x.size(-1) != self.n_embd:
            raise ValueError(
                f"Expected hidden size {self.n_embd}, got {x.size(-1)}"
            )

        batch_size, seq_len, _ = x.shape
        grouped = self.norm(x).reshape(
            batch_size, seq_len, self.n_group, self.d_group
        )

        intra_output = self._intra_mix(grouped)
        intra_gate = torch.sigmoid(self.intra_gate_logits).view(
            1, 1, self.n_group, self.d_group
        )
        mixed_input = grouped + intra_gate * intra_output

        update = self._cross_mix(mixed_input).reshape(
            batch_size, seq_len, self.n_embd
        )
        return x + self.drop(update)


class TransformerFFNBlock(nn.Module):
    def __init__(
        self,
        n_embd: int,
        n_head: int,

        attn_dropout: float = 0.0,
        mlp_dropout: float = 0.0,
        use_time_rope: bool = False,
        use_rbf_bias: bool = False,
        n_rbf_bases: int = 16,
        attention_type: str = STANDARD_ATTENTION,
        depth: int = 0,
    ):
        super().__init__()
        self.attn = build_attention(
            attention_type,
            n_embd=n_embd,
            n_head=n_head,
            n_rbf_bases=n_rbf_bases,
            dropout=attn_dropout,
            use_time_rope=use_time_rope,
            use_rbf_bias=use_rbf_bias,
            depth=depth,
        )
        self.mlp = SwiGLU(n_embd=n_embd, dropout=mlp_dropout)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(
        self,
        x: torch.Tensor,
        rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
        rbf_cache: torch.Tensor | None = None,
        attn_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        x = x + self.attn(self.ln1(x), rope_cache, rbf_cache, attn_mask)
        x = x + self.mlp(self.ln2(x))
        return x


class TrajMixerBlock(nn.Module):
    def __init__(
        self,
        n_embd: int,
        n_head: int,
        attn_dropout: float = 0.0,
        mlp_dropout: float = 0.0,
        use_time_rope: bool = False,
        use_rbf_bias: bool = False,
        n_rbf_bases: int = 16,
        attention_type: str = STANDARD_ATTENTION,
        depth: int = 0,
    ):
        super().__init__()
        self.attn = build_attention(
            attention_type,
            n_embd=n_embd,
            n_head=n_head,
            n_rbf_bases=n_rbf_bases,
            dropout=attn_dropout,
            use_time_rope=use_time_rope,
            use_rbf_bias=use_rbf_bias,
            depth=depth,
        )
        self.mlp = TrajMixer(
            n_embd=n_embd,
            n_head=n_head,
            dropout=mlp_dropout,
        )
        self.ln1 = nn.LayerNorm(n_embd)

    def forward(
        self,
        x: torch.Tensor,
        rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
        rbf_cache: torch.Tensor | None = None,
        attn_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        x = x + self.attn(self.ln1(x), rope_cache, rbf_cache, attn_mask)
        return self.mlp(x)


def build_backbone_block(
    model_architecture: str,
    *,
    n_embd: int,
    n_head: int,
    attn_dropout: float = 0.0,
    mlp_dropout: float = 0.0,
    use_time_rope: bool = False,
    use_rbf_bias: bool = False,
    n_rbf_bases: int = 16,
    attention_type: str = STANDARD_ATTENTION,
    depth: int = 0,
) -> nn.Module:
    """Build one history block for a supported model architecture."""
    architecture = resolve_model_architecture(model_architecture)
    block_class: type[nn.Module]
    if architecture == TRANSFORMER_FFN_ARCHITECTURE:
        block_class = TransformerFFNBlock
    elif architecture == TRAJ_MIXER_ARCHITECTURE:
        block_class = TrajMixerBlock
    else:  # pragma: no cover - guarded by resolve_model_architecture.
        raise ValueError(f"Unsupported model architecture: {architecture!r}")
    return block_class(
        n_embd=n_embd,
        n_head=n_head,
        attn_dropout=attn_dropout,
        mlp_dropout=mlp_dropout,
        use_time_rope=use_time_rope,
        use_rbf_bias=use_rbf_bias,
        n_rbf_bases=n_rbf_bases,
        attention_type=attention_type,
        depth=depth,
    )


class TokenAutoDiscretization(nn.Module):
    def __init__(
        self,
        n_cont_types: int,
        n_bins: int,
        n_embd: int,
    ):
        super().__init__()
        if n_cont_types <= 0:
            raise ValueError(f"n_cont_types must be > 0, got {n_cont_types}")
        if n_bins <= 1:
            raise ValueError(f"n_bins must be > 1, got {n_bins}")
        if n_embd <= 0:
            raise ValueError(f"n_embd must be > 0, got {n_embd}")

        self.n_cont_types = n_cont_types
        self.n_bins = n_bins
        self.n_embd = n_embd
        self.weight = nn.Parameter(torch.empty(n_cont_types, n_bins))
        self.bias = nn.Parameter(torch.empty(n_cont_types, n_bins))
        self.bin_emb = nn.Parameter(torch.empty(n_cont_types, n_bins, n_embd))
        self.reset_parameters()

    def reset_parameters(self) -> None:
        nn.init.normal_(self.weight, mean=0.0, std=0.02)
        nn.init.zeros_(self.bias)
        nn.init.normal_(self.bin_emb, mean=0.0, std=0.02)

    def forward(
        self,
        cont_type_idx: torch.LongTensor,  # (N,)
        value: torch.Tensor,              # (N,)
    ) -> torch.Tensor:
        if cont_type_idx.dim() != 1:
            raise ValueError(
                f"cont_type_idx must be 1D, got {tuple(cont_type_idx.shape)}"
            )
        if value.dim() != 1:
            raise ValueError(f"value must be 1D, got {tuple(value.shape)}")
        if cont_type_idx.numel() != value.numel():
            raise ValueError("cont_type_idx and value must have the same length")

        w = self.weight[cont_type_idx]      # (N, n_bins)
        b = self.bias[cont_type_idx]        # (N, n_bins)
        e = self.bin_emb[cont_type_idx]     # (N, n_bins, D)
        logits = value.unsqueeze(-1) * w + b
        probs = torch.softmax(logits, dim=-1)
        return torch.einsum("nb,nbd->nd", probs, e)



class AgeSinusoidalEncoding(nn.Module):

    def __init__(self, embedding_dim: int):

        super().__init__()
        if embedding_dim % 2 != 0:
            raise ValueError(
                f"Embedding dimension must be an even number, but got {embedding_dim}")

        self.embedding_dim = embedding_dim

        i = torch.arange(0, self.embedding_dim, 2, dtype=torch.float32)
        divisor = torch.pow(10000, i / self.embedding_dim)
        self.register_buffer('divisor', divisor)
        self.linear = nn.Linear(embedding_dim, embedding_dim, bias=False)

    def forward(self, t: torch.Tensor) -> torch.Tensor:

        t_years = t
        # Broadcast (B, L, 1) against (1, 1, D/2) to get (B, L, D/2)
        args = t_years.unsqueeze(-1) / self.divisor.view(1, 1, -1)
        # Interleave cos and sin along the last dimension
        output = torch.zeros(t.shape[0], t.shape[1],
                             self.embedding_dim, device=t.device)
        output[:, :, 0::2] = torch.cos(args)
        output[:, :, 1::2] = torch.sin(args)
        output = self.linear(output)
        return output