#!/usr/bin/env python3
"""Inspectable transformer parts implemented without ``nn.Transformer``.

Mask convention
---------------
All boolean masks in this file use ``True`` for an allowed token or attention
connection.  This is intentionally stated here because some PyTorch APIs use
the opposite convention.  Attention raises ``ValueError`` instead of silently
running softmax when a query has no allowed key.

Important shapes use these names:

* ``B``: batch size
* ``T``: sequence length
* ``D``: model width (``d_model``)
* ``H``: number of attention heads
* ``Dh``: width of one head (``D // H``)
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Sequence

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


def _check_valid_mask(mask: torch.Tensor, shape: tuple[int, int], name: str) -> None:
    if mask.dtype != torch.bool:
        raise TypeError(f"{name} must have dtype torch.bool (True means valid)")
    if tuple(mask.shape) != shape:
        raise ValueError(f"{name} must have shape {shape}, got {tuple(mask.shape)}")


def rotate_pairs(x: torch.Tensor) -> torch.Tensor:
    """Rotate every adjacent pair ``(a, b)`` to ``(-b, a)``."""
    if x.shape[-1] % 2:
        raise ValueError("RoPE requires an even final dimension")
    pairs = x.reshape(*x.shape[:-1], -1, 2)
    first, second = pairs.unbind(dim=-1)
    return torch.stack((-second, first), dim=-1).flatten(start_dim=-2)


def rope_sin_cos(
    length: int,
    width: int,
    *,
    offset: int = 0,
    base: float = 10_000.0,
    device: torch.device | None = None,
    dtype: torch.dtype = torch.float32,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return RoPE sine/cosine tables with shape ``(length, width)``."""
    if width % 2:
        raise ValueError("RoPE width must be even")
    if length < 0 or offset < 0:
        raise ValueError("RoPE length and offset must be non-negative")

    # Compute angles in float32 for stable CPU behavior, then match x's dtype.
    positions = torch.arange(offset, offset + length, device=device, dtype=torch.float32)
    pair_indices = torch.arange(0, width, 2, device=device, dtype=torch.float32)
    inverse_frequencies = base ** (-pair_indices / width)
    pair_angles = positions[:, None] * inverse_frequencies[None, :]
    angles = torch.repeat_interleave(pair_angles, repeats=2, dim=-1)
    return angles.sin().to(dtype=dtype), angles.cos().to(dtype=dtype)


def apply_rope(x: torch.Tensor, *, offset: int = 0, base: float = 10_000.0) -> torch.Tensor:
    """Apply rotary position embedding to ``(..., T, Dh)`` with an offset."""
    length, width = x.shape[-2:]
    sin, cos = rope_sin_cos(
        length,
        width,
        offset=offset,
        base=base,
        device=x.device,
        dtype=x.dtype,
    )
    # Add leading singleton dimensions so (T, Dh) broadcasts to (B, H, T, Dh).
    while sin.ndim < x.ndim:
        sin = sin.unsqueeze(0)
        cos = cos.unsqueeze(0)
    return x * cos + rotate_pairs(x) * sin


class RotaryEmbedding(nn.Module):
    """Small module wrapper around :func:`apply_rope`."""

    def __init__(self, base: float = 10_000.0) -> None:
        super().__init__()
        self.base = base

    def forward(self, x: torch.Tensor, offset: int = 0) -> torch.Tensor:
        return apply_rope(x, offset=offset, base=self.base)


@dataclass(frozen=True)
class KVCache:
    """Keys/values for one decoder layer, both shaped ``(B, H, T, Dh)``."""

    key: torch.Tensor
    value: torch.Tensor
    key_padding_mask: torch.Tensor  # (B, T), True for a real token

    def __post_init__(self) -> None:
        if self.key.ndim != 4 or self.value.ndim != 4:
            raise ValueError("cached key and value must have shape (B, H, T, Dh)")
        if self.key.shape != self.value.shape:
            raise ValueError("cached key and value shapes must match")
        expected = (self.key.shape[0], self.key.shape[2])
        _check_valid_mask(self.key_padding_mask, expected, "cached key_padding_mask")

    @property
    def length(self) -> int:
        return self.key.shape[2]


class ManualLayerNorm(nn.Module):
    """Layer normalization written from mean and variance for inspection."""

    def __init__(self, width: int, eps: float = 1e-5) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(width))
        self.bias = nn.Parameter(torch.zeros(width))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mean = x.mean(dim=-1, keepdim=True)
        variance = ((x - mean) ** 2).mean(dim=-1, keepdim=True)
        normalized = (x - mean) * torch.rsqrt(variance + self.eps)
        return normalized * self.weight + self.bias


class ManualMultiHeadAttention(nn.Module):
    """Multi-head attention with explicit score, scale, mask, and softmax.

    ``query`` has shape ``(B, Tq, D)``. ``key_value`` has shape ``(B, Tk, D)``
    and defaults to ``query`` for self-attention.  Normally the return value is
    a tensor ``(B, Tq, D)``.  With ``use_cache=True`` the return value is
    ``(output, new_cache)``.
    """

    def __init__(
        self,
        d_model: int,
        num_heads: int,
        *,
        dropout: float = 0.0,
        use_rope: bool = False,
        bias: bool = True,
    ) -> None:
        super().__init__()
        if d_model % num_heads:
            raise ValueError("d_model must be divisible by num_heads")
        head_dim = d_model // num_heads
        if use_rope and head_dim % 2:
            raise ValueError("RoPE requires an even head dimension")
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.scale = head_dim**-0.5
        self.use_rope = use_rope
        self.q_proj = nn.Linear(d_model, d_model, bias=bias)
        self.k_proj = nn.Linear(d_model, d_model, bias=bias)
        self.v_proj = nn.Linear(d_model, d_model, bias=bias)
        self.out_proj = nn.Linear(d_model, d_model, bias=bias)
        self.attention_dropout = nn.Dropout(dropout)
        self.output_dropout = nn.Dropout(dropout)
        self.rope = RotaryEmbedding()

    def split_heads(self, x: torch.Tensor) -> torch.Tensor:
        """Convert ``(B, T, D)`` to ``(B, H, T, Dh)``."""
        batch, length, width = x.shape
        if width != self.d_model:
            raise ValueError(f"last dimension must be d_model={self.d_model}")
        return x.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2)

    def merge_heads(self, x: torch.Tensor) -> torch.Tensor:
        """Convert ``(B, H, T, Dh)`` back to ``(B, T, D)``."""
        batch, heads, length, head_dim = x.shape
        if heads != self.num_heads or head_dim != self.head_dim:
            raise ValueError("head dimensions do not match this attention module")
        return x.transpose(1, 2).contiguous().reshape(batch, length, self.d_model)

    def _expand_attention_mask(
        self,
        mask: torch.Tensor,
        batch: int,
        query_length: int,
        key_length: int,
    ) -> torch.Tensor:
        if mask.dtype != torch.bool:
            raise TypeError("attention_mask must be bool (True means allowed)")
        if mask.ndim == 2:
            if tuple(mask.shape) != (query_length, key_length):
                raise ValueError("2D attention_mask must have shape (Tq, Tk)")
            return mask[None, None, :, :]
        if mask.ndim == 3:
            if tuple(mask.shape) != (batch, query_length, key_length):
                raise ValueError("3D attention_mask must have shape (B, Tq, Tk)")
            return mask[:, None, :, :]
        if mask.ndim == 4 and mask.shape[0] in (1, batch) and mask.shape[1] in (
            1,
            self.num_heads,
        ):
            if tuple(mask.shape[-2:]) != (query_length, key_length):
                raise ValueError("4D attention_mask has incorrect sequence dimensions")
            return mask
        raise ValueError("attention_mask must have 2, 3, or 4 dimensions")

    def forward(
        self,
        query: torch.Tensor,
        key_value: torch.Tensor | None = None,
        *,
        attention_mask: torch.Tensor | None = None,
        key_padding_mask: torch.Tensor | None = None,
        query_padding_mask: torch.Tensor | None = None,
        causal: bool = False,
        cache: KVCache | None = None,
        use_cache: bool = False,
        query_position_offset: int | None = None,
        key_position_offset: int = 0,
    ) -> torch.Tensor | tuple[torch.Tensor, KVCache]:
        if query.ndim != 3 or query.shape[-1] != self.d_model:
            raise ValueError("query must have shape (B, Tq, d_model)")
        if query.shape[1] == 0:
            raise ValueError("attention needs at least one query token")
        if cache is not None and not use_cache:
            raise ValueError("cache input requires use_cache=True so the updated cache is returned")
        if use_cache:
            expected_offset = cache.length if cache is not None else 0
            if key_position_offset != 0 or query_position_offset not in (None, expected_offset):
                raise ValueError("cached positions must start at zero and continue from cache length")
        if cache is not None and key_value is not None and key_value is not query:
            raise ValueError("KV caching is only supported for self-attention")
        if use_cache and key_value is not None and key_value is not query:
            raise ValueError("KV caching is only supported for self-attention")
        if use_cache and self.training and (
            self.attention_dropout.p > 0.0 or self.output_dropout.p > 0.0
        ):
            raise ValueError("KV caching with dropout is inference-only; call eval() first")

        is_self_attention = key_value is None or key_value is query
        source = query if key_value is None else key_value
        if source.ndim != 3 or source.shape[0] != query.shape[0] or source.shape[-1] != self.d_model:
            raise ValueError("key_value must have shape (B, Tk, d_model)")

        batch, query_length, _ = query.shape
        current_key_length = source.shape[1]
        if key_padding_mask is None:
            current_valid = torch.ones(
                (batch, current_key_length), dtype=torch.bool, device=query.device
            )
        else:
            _check_valid_mask(
                key_padding_mask, (batch, current_key_length), "key_padding_mask"
            )
            current_valid = key_padding_mask
        if query_padding_mask is not None:
            _check_valid_mask(
                query_padding_mask, (batch, query_length), "query_padding_mask"
            )

        q = self.split_heads(self.q_proj(query))
        new_k = self.split_heads(self.k_proj(source))
        new_v = self.split_heads(self.v_proj(source))

        inferred_offset = cache.length if cache is not None else (
            key_position_offset if is_self_attention else 0
        )
        q_offset = inferred_offset if query_position_offset is None else query_position_offset
        if q_offset < 0 or key_position_offset < 0:
            raise ValueError("position offsets must be non-negative")
        if self.use_rope:
            q = self.rope(q, offset=q_offset)
            new_k_offset = inferred_offset if cache is not None else key_position_offset
            new_k = self.rope(new_k, offset=new_k_offset)

        if cache is not None:
            if cache.key.shape[:2] != (batch, self.num_heads):
                raise ValueError("cache batch/head dimensions do not match")
            if cache.key.shape[-1] != self.head_dim:
                raise ValueError("cache head width does not match")
            k = torch.cat((cache.key, new_k), dim=2)
            v = torch.cat((cache.value, new_v), dim=2)
            valid_keys = torch.cat((cache.key_padding_mask, current_valid), dim=1)
        else:
            k, v, valid_keys = new_k, new_v, current_valid

        key_length = k.shape[2]
        allowed = valid_keys[:, None, None, :]  # (B, 1, 1, Tk)
        if causal:
            query_positions = torch.arange(
                q_offset, q_offset + query_length, device=query.device
            )
            first_key_position = 0 if cache is not None else key_position_offset
            key_positions = torch.arange(
                first_key_position, first_key_position + key_length, device=query.device
            )
            causal_mask = key_positions[None, :] <= query_positions[:, None]
            allowed = allowed & causal_mask[None, None, :, :]
        if attention_mask is not None:
            expanded = self._expand_attention_mask(
                attention_mask, batch, query_length, key_length
            )
            allowed = allowed & expanded

        # Expand before checking so each batch/head/query row is checked.
        allowed = allowed.expand(batch, self.num_heads, query_length, key_length)
        if not bool(allowed.any(dim=-1).all()):
            raise ValueError("attention mask contains an all-masked query row")

        # (B,H,Tq,Dh) @ (B,H,Dh,Tk) -> one score per query/key pair.
        scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        scores = scores.masked_fill(~allowed, torch.finfo(scores.dtype).min)
        weights = torch.softmax(scores, dim=-1)
        weights = self.attention_dropout(weights)
        mixed_values = torch.matmul(weights, v)  # (B, H, Tq, Dh)
        output = self.out_proj(self.merge_heads(mixed_values))
        output = self.output_dropout(output)
        if query_padding_mask is not None:
            output = output * query_padding_mask.unsqueeze(-1).to(output.dtype)

        if use_cache:
            return output, KVCache(k, v, valid_keys)
        return output


class FeedForward(nn.Module):
    """Position-wise ``D -> hidden -> D`` network used by each block."""

    def __init__(self, d_model: int, hidden_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.input_projection = nn.Linear(d_model, hidden_dim)
        self.output_projection = nn.Linear(hidden_dim, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.dropout(self.output_projection(self.dropout(F.gelu(self.input_projection(x)))))


class EncoderBlock(nn.Module):
    """Pre-norm encoder block: non-causal self-attention, then FFN."""

    def __init__(self, d_model: int, num_heads: int, hidden_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.attention_norm = ManualLayerNorm(d_model)
        self.self_attention = ManualMultiHeadAttention(
            d_model, num_heads, dropout=dropout, use_rope=True
        )
        self.ffn_norm = ManualLayerNorm(d_model)
        self.feed_forward = FeedForward(d_model, hidden_dim, dropout)

    def forward(self, x: torch.Tensor, valid_tokens: torch.Tensor) -> torch.Tensor:
        attention_input = self.attention_norm(x)
        x = x + self.self_attention(
            attention_input,
            key_padding_mask=valid_tokens,
            query_padding_mask=valid_tokens,
        )
        x = x + self.feed_forward(self.ffn_norm(x))
        return x * valid_tokens.unsqueeze(-1).to(x.dtype)


class DecoderBlock(nn.Module):
    """Pre-norm decoder block with causal self-attention and cross-attention."""

    def __init__(self, d_model: int, num_heads: int, hidden_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.self_norm = ManualLayerNorm(d_model)
        self.self_attention = ManualMultiHeadAttention(
            d_model, num_heads, dropout=dropout, use_rope=True
        )
        self.cross_norm = ManualLayerNorm(d_model)
        self.cross_attention = ManualMultiHeadAttention(
            d_model, num_heads, dropout=dropout, use_rope=True
        )
        self.ffn_norm = ManualLayerNorm(d_model)
        self.feed_forward = FeedForward(d_model, hidden_dim, dropout)

    def forward(
        self,
        x: torch.Tensor,
        memory: torch.Tensor,
        *,
        target_valid_tokens: torch.Tensor,
        memory_valid_tokens: torch.Tensor,
    ) -> torch.Tensor:
        normalized = self.self_norm(x)
        x = x + self.self_attention(
            normalized,
            key_padding_mask=target_valid_tokens,
            query_padding_mask=target_valid_tokens,
            causal=True,
        )
        normalized = self.cross_norm(x)
        x = x + self.cross_attention(
            normalized,
            memory,
            key_padding_mask=memory_valid_tokens,
            query_padding_mask=target_valid_tokens,
        )
        x = x + self.feed_forward(self.ffn_norm(x))
        return x * target_valid_tokens.unsqueeze(-1).to(x.dtype)


class DecoderOnlyBlock(nn.Module):
    """Pre-norm causal block that can append to a per-layer KV cache."""

    def __init__(self, d_model: int, num_heads: int, hidden_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.attention_norm = ManualLayerNorm(d_model)
        self.self_attention = ManualMultiHeadAttention(
            d_model, num_heads, dropout=dropout, use_rope=True
        )
        self.ffn_norm = ManualLayerNorm(d_model)
        self.feed_forward = FeedForward(d_model, hidden_dim, dropout)

    def forward(
        self,
        x: torch.Tensor,
        valid_tokens: torch.Tensor,
        *,
        cache: KVCache | None = None,
        use_cache: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, KVCache]:
        attention_result = self.self_attention(
            self.attention_norm(x),
            key_padding_mask=valid_tokens,
            query_padding_mask=valid_tokens,
            causal=True,
            cache=cache,
            use_cache=use_cache,
        )
        if use_cache:
            attention_output, new_cache = attention_result
        else:
            attention_output = attention_result
        x = x + attention_output
        x = x + self.feed_forward(self.ffn_norm(x))
        x = x * valid_tokens.unsqueeze(-1).to(x.dtype)
        if use_cache:
            return x, new_cache
        return x


class TransformerEncoder(nn.Module):
    """Stack of non-causal encoder blocks operating on embedded vectors."""

    def __init__(
        self, d_model: int, num_heads: int, hidden_dim: int, num_layers: int, dropout: float = 0.0
    ) -> None:
        super().__init__()
        self.layers = nn.ModuleList(
            [EncoderBlock(d_model, num_heads, hidden_dim, dropout) for _ in range(num_layers)]
        )
        self.final_norm = ManualLayerNorm(d_model)

    def forward(self, x: torch.Tensor, valid_tokens: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            x = layer(x, valid_tokens)
        return self.final_norm(x) * valid_tokens.unsqueeze(-1).to(x.dtype)


class TransformerDecoder(nn.Module):
    """Stack of causal decoder blocks operating on embedded vectors."""

    def __init__(
        self, d_model: int, num_heads: int, hidden_dim: int, num_layers: int, dropout: float = 0.0
    ) -> None:
        super().__init__()
        self.layers = nn.ModuleList(
            [DecoderBlock(d_model, num_heads, hidden_dim, dropout) for _ in range(num_layers)]
        )
        self.final_norm = ManualLayerNorm(d_model)

    def forward(
        self,
        x: torch.Tensor,
        memory: torch.Tensor,
        *,
        target_valid_tokens: torch.Tensor,
        memory_valid_tokens: torch.Tensor,
    ) -> torch.Tensor:
        for layer in self.layers:
            x = layer(
                x,
                memory,
                target_valid_tokens=target_valid_tokens,
                memory_valid_tokens=memory_valid_tokens,
            )
        return self.final_norm(x) * target_valid_tokens.unsqueeze(-1).to(x.dtype)


class Seq2SeqTransformer(nn.Module):
    """Complete encoder-decoder transformer for token IDs."""

    def __init__(
        self,
        vocab_size: int,
        *,
        d_model: int = 64,
        num_heads: int = 4,
        hidden_dim: int = 128,
        num_layers: int = 2,
        dropout: float = 0.0,
        padding_id: int = 0,
    ) -> None:
        super().__init__()
        self.padding_id = padding_id
        self.d_model = d_model
        self.source_embedding = nn.Embedding(vocab_size, d_model, padding_idx=padding_id)
        self.target_embedding = nn.Embedding(vocab_size, d_model, padding_idx=padding_id)
        self.encoder = TransformerEncoder(d_model, num_heads, hidden_dim, num_layers, dropout)
        self.decoder = TransformerDecoder(d_model, num_heads, hidden_dim, num_layers, dropout)
        self.output_projection = nn.Linear(d_model, vocab_size, bias=False)

    def encode(
        self, source_ids: torch.Tensor, source_valid: torch.Tensor | None = None
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if source_valid is None:
            source_valid = source_ids.ne(self.padding_id)
        _check_valid_mask(source_valid, tuple(source_ids.shape), "source_valid")
        if not bool(source_valid.any(dim=1).all()):
            raise ValueError("each source sequence needs at least one valid token")
        embedded = self.source_embedding(source_ids) * math.sqrt(self.d_model)
        return self.encoder(embedded, source_valid), source_valid

    def decode(
        self,
        target_ids: torch.Tensor,
        memory: torch.Tensor,
        *,
        target_valid: torch.Tensor | None = None,
        memory_valid: torch.Tensor,
    ) -> torch.Tensor:
        if target_valid is None:
            target_valid = target_ids.ne(self.padding_id)
        embedded = self.target_embedding(target_ids) * math.sqrt(self.d_model)
        hidden = self.decoder(
            embedded,
            memory,
            target_valid_tokens=target_valid,
            memory_valid_tokens=memory_valid,
        )
        return self.output_projection(hidden)

    def forward(
        self,
        source_ids: torch.Tensor,
        target_input_ids: torch.Tensor,
        *,
        source_valid: torch.Tensor | None = None,
        target_valid: torch.Tensor | None = None,
    ) -> torch.Tensor:
        memory, source_valid = self.encode(source_ids, source_valid)
        return self.decode(
            target_input_ids,
            memory,
            target_valid=target_valid,
            memory_valid=source_valid,
        )

    @torch.no_grad()
    def generate(
        self,
        source_ids: torch.Tensor,
        *,
        bos_id: int,
        eos_id: int,
        max_new_tokens: int,
        source_valid: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Greedy decode and return generated IDs, including BOS and possible EOS."""
        self.eval()
        memory, source_valid = self.encode(source_ids, source_valid)
        generated = torch.full(
            (source_ids.shape[0], 1), bos_id, dtype=torch.long, device=source_ids.device
        )
        finished = torch.zeros(source_ids.shape[0], dtype=torch.bool, device=source_ids.device)
        for _ in range(max_new_tokens):
            logits = self.decode(generated, memory, memory_valid=source_valid)
            next_ids = logits[:, -1].argmax(dim=-1)
            next_ids = torch.where(finished, torch.full_like(next_ids, eos_id), next_ids)
            generated = torch.cat((generated, next_ids[:, None]), dim=1)
            finished |= next_ids.eq(eos_id)
            if bool(finished.all()):
                break
        return generated


class DecoderOnlyTransformer(nn.Module):
    """Decoder-only language model with one :class:`KVCache` per layer."""

    def __init__(
        self,
        vocab_size: int,
        *,
        d_model: int = 64,
        num_heads: int = 4,
        hidden_dim: int = 128,
        num_layers: int = 2,
        dropout: float = 0.0,
        padding_id: int = 0,
    ) -> None:
        super().__init__()
        self.padding_id = padding_id
        self.d_model = d_model
        self.token_embedding = nn.Embedding(vocab_size, d_model, padding_idx=padding_id)
        self.layers = nn.ModuleList(
            [DecoderOnlyBlock(d_model, num_heads, hidden_dim, dropout) for _ in range(num_layers)]
        )
        self.final_norm = ManualLayerNorm(d_model)

    def forward(
        self,
        token_ids: torch.Tensor,
        *,
        valid_tokens: torch.Tensor | None = None,
        caches: Sequence[KVCache | None] | None = None,
        use_cache: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, list[KVCache]]:
        if token_ids.ndim != 2:
            raise ValueError("token_ids must have shape (B, T)")
        if valid_tokens is None:
            valid_tokens = token_ids.ne(self.padding_id)
        _check_valid_mask(valid_tokens, tuple(token_ids.shape), "valid_tokens")
        if not bool(valid_tokens.any(dim=1).all()):
            raise ValueError("each input sequence needs at least one valid token")
        if caches is None:
            caches = [None] * len(self.layers)
        if len(caches) != len(self.layers):
            raise ValueError("provide exactly one cache entry per layer")
        initialized = [cache is not None for cache in caches]
        if any(initialized) and not all(initialized):
            raise ValueError("layer caches must be either all initialized or all empty")
        lengths = {cache.length for cache in caches if cache is not None}
        if len(lengths) > 1:
            raise ValueError("all layer caches must have the same sequence length")

        x = self.token_embedding(token_ids) * math.sqrt(self.d_model)
        new_caches: list[KVCache] = []
        for layer, cache in zip(self.layers, caches):
            result = layer(x, valid_tokens, cache=cache, use_cache=use_cache)
            if use_cache:
                x, new_cache = result
                new_caches.append(new_cache)
            else:
                x = result
        x = self.final_norm(x) * valid_tokens.unsqueeze(-1).to(x.dtype)
        # Tied output weights avoid a separate, hidden language-model projection.
        logits = F.linear(x, self.token_embedding.weight)
        if use_cache:
            return logits, new_caches
        return logits

    @torch.no_grad()
    def generate_cached(self, prefix_ids: torch.Tensor, max_new_tokens: int) -> torch.Tensor:
        """Prefill once, then decode one token at a time using the KV caches."""
        self.eval()
        generated = prefix_ids
        logits, caches = self(prefix_ids, use_cache=True)
        for _ in range(max_new_tokens):
            next_id = logits[:, -1].argmax(dim=-1, keepdim=True)
            generated = torch.cat((generated, next_id), dim=1)
            logits, caches = self(next_id, caches=caches, use_cache=True)
        return generated


__all__ = [
    "DecoderBlock",
    "DecoderOnlyBlock",
    "DecoderOnlyTransformer",
    "EncoderBlock",
    "FeedForward",
    "KVCache",
    "ManualLayerNorm",
    "ManualMultiHeadAttention",
    "RotaryEmbedding",
    "Seq2SeqTransformer",
    "TransformerDecoder",
    "TransformerEncoder",
    "apply_rope",
    "rope_sin_cos",
    "rotate_pairs",
]