# bclm/model.py
"""
BCLM-1 Model Architecture

Hybrid language model combining oscillator-based sequence mixing,
sparse sliding-window / global causal attention (GQA), and gated
feed-forward networks.  Internal layout: [B, T, C] throughout.
"""
from __future__ import annotations

import inspect
import math
import warnings
from functools import lru_cache
from typing import Optional, Set, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint as grad_checkpoint

from .config import BCLM1ModelConfig


# =============================================================================
# Optional xFormers
# =============================================================================

try:
    import xformers.ops as xops
    from xformers.ops import fmha

    _XFORMERS_AVAILABLE = True
except Exception:
    xops, fmha = None, None  # type: ignore[assignment]
    _XFORMERS_AVAILABLE = False


def _can_use_xformers_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
) -> bool:
    if not _XFORMERS_AVAILABLE:
        return False
    if not (q.is_cuda and k.is_cuda and v.is_cuda):
        return False
    if q.dtype not in (torch.float16, torch.bfloat16, torch.float32):
        return False
    if k.dtype != q.dtype or v.dtype != q.dtype:
        return False
    if q.stride(-1) != 1 or k.stride(-1) != 1 or v.stride(-1) != 1:
        return False
    if q.size(-1) != k.size(-1) or q.size(-1) != v.size(-1):
        return False
    if q.size(-2) % k.size(-2) != 0:
        return False
    if k.size(-2) != v.size(-2):
        return False
    return True


@lru_cache(maxsize=None)
def _sdpa_supports_enable_gqa() -> bool:
    try:
        sig = inspect.signature(F.scaled_dot_product_attention)
        return "enable_gqa" in sig.parameters
    except Exception:
        return False


# =============================================================================
# Numeric Helpers
# =============================================================================


@lru_cache(maxsize=64)
def _next_fast_fft_size(n: int) -> int:
    if n <= 1:
        return 1
    return 1 << (n - 1).bit_length()


def _check_tensor_core_alignment(dim: int, label: str) -> None:
    if dim % 64 != 0:
        warnings.warn(
            f"{label}={dim} is not a multiple of 64; "
            "Tensor-Core utilisation will be reduced.",
            stacklevel=3,
        )


# =============================================================================
# GQA Head-Expansion Helpers
# =============================================================================


def _expand_kv_for_xformers(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    hq, hkv = q.size(-2), k.size(-2)
    if hkv == hq:
        return k, v
    rep = hq // hkv
    return (
        k.repeat_interleave(rep, dim=-2).contiguous(),
        v.repeat_interleave(rep, dim=-2).contiguous(),
    )


def _repeat_kv_heads(x: torch.Tensor, n_heads: int) -> torch.Tensor:
    b, hkv, t, d = x.shape
    if hkv == n_heads:
        return x
    return x.repeat_interleave(n_heads // hkv, dim=1)


# =============================================================================
# JIT-Scripted Oscillator Kernel Builder
# =============================================================================


@torch.jit.script
def _build_osc_kernel(
    t: torch.Tensor,
    log_alpha: torch.Tensor,
    log_omega: torch.Tensor,
    phi: torch.Tensor,
    log_amp: torch.Tensor,
    log_beta: torch.Tensor,
    log_amp2: torch.Tensor,
    sign2: torch.Tensor,
    clamp_min: float,
) -> torch.Tensor:
    """
    Build the causal impulse-response kernel from oscillator parameters.

    Args:
        t:         [T]    time indices.
        log_alpha: [C, P] log decay rates (oscillator).
        log_omega: [C, P] log angular frequencies.
        phi:       [C, P] phase offsets.
        log_amp:   [C, P] log amplitudes (oscillator).
        log_beta:  [C, R] log decay rates (real exponential).
        log_amp2:  [C, R] log amplitudes (real exponential).
        sign2:     [C, R] sign controls via tanh.
        clamp_min: Lower bound for decay rates after exp.

    Returns:
        [C, T] combined causal kernel.
    """
    alpha = torch.exp(log_alpha).clamp_min(clamp_min)
    omega = torch.exp(log_omega)
    amp = torch.exp(log_amp)

    at = alpha.unsqueeze(-1) * t
    phase = omega.unsqueeze(-1) * t + phi.unsqueeze(-1)
    k_osc = (amp.unsqueeze(-1) * torch.exp(-at) * torch.cos(phase)).sum(1)

    beta = torch.exp(log_beta).clamp_min(clamp_min)
    coeff = torch.tanh(sign2) * torch.exp(log_amp2)
    bt = beta.unsqueeze(-1) * t
    k_exp = (coeff.unsqueeze(-1) * torch.exp(-bt)).sum(1)

    return k_osc + k_exp


# =============================================================================
# Normalization
# =============================================================================


class RMSNorm(nn.Module):
    """Root-Mean-Square layer normalisation on the last dimension."""

    __constants__ = ["dim", "eps"]

    def __init__(self, dim: int, eps: float = 1e-8) -> None:
        super().__init__()
        self.dim = dim
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.rms_norm(x, (self.dim,), weight=self.weight, eps=self.eps)


# =============================================================================
# Causal Convolution
# =============================================================================


class CausalConv(nn.Module):
    """
    Causal 1-D convolution.

    Accepts and returns [B, T, C]; transposes internally for Conv1d.
    """

    def __init__(self, dim: int, kernel_size: int = 4, dropout: float = 0.0) -> None:
        super().__init__()
        self.kernel_size = kernel_size
        self.conv = nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=0, bias=False)
        nn.init.zeros_(self.conv.weight)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        u = x.transpose(1, 2).contiguous()
        u = self.conv(F.pad(u, (self.kernel_size - 1, 0)))
        return self.dropout(u.transpose(1, 2))


# =============================================================================
# Oscillator Filter
# =============================================================================


class CausalOscillatorFilter(nn.Module):
    """
    Learnable causal filter via damped oscillators + exponential decays.

    Performs FFT convolution in [B, C, T]; accepts and returns [B, T, C].
    The eval-mode kernel-FFT cache is invalidated when switching back to
    training.
    """

    def __init__(
        self,
        dim: int,
        n_pairs: int,
        n_real: int,
        clamp_min_decay: float,
        max_len: int,
    ) -> None:
        super().__init__()
        self.clamp_min_decay = float(clamp_min_decay)

        self.register_buffer(
            "_t",
            torch.arange(max_len, dtype=torch.float32),
            persistent=False,
        )

        self.log_alpha = nn.Parameter(torch.randn(dim, n_pairs) * 0.1 - 2.0)
        self.log_omega = nn.Parameter(torch.randn(dim, n_pairs) * 0.1 - 1.0)
        self.phi = nn.Parameter(torch.randn(dim, n_pairs) * 0.5)
        self.log_amp = nn.Parameter(torch.randn(dim, n_pairs) * 0.1 - 1.0)

        self.log_beta = nn.Parameter(torch.randn(dim, n_real) * 0.1 - 2.0)
        self.log_amp2 = nn.Parameter(torch.randn(dim, n_real) * 0.1 - 1.0)
        self.sign2 = nn.Parameter(torch.zeros(dim, n_real))

        self._cached_k_f: Optional[torch.Tensor] = None
        self._cached_n_fft: int = -1

    def _invalidate_cache(self) -> None:
        self._cached_k_f = None
        self._cached_n_fft = -1

    def train(self, mode: bool = True):
        super().train(mode)
        if mode:
            self._invalidate_cache()
        return self

    def _get_time(self, seqlen: int, device: torch.device) -> torch.Tensor:
        if seqlen <= self._t.size(0):
            return self._t[:seqlen]
        return torch.arange(seqlen, device=device, dtype=torch.float32)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        seqlen = x.size(1)
        x_dtype = x.dtype

        x_ct = x.transpose(1, 2).contiguous().float()
        n_fft = _next_fast_fft_size(2 * seqlen)

        if self.training:
            t = self._get_time(seqlen, x.device)
            k = _build_osc_kernel(
                t,
                self.log_alpha, self.log_omega, self.phi, self.log_amp,
                self.log_beta, self.log_amp2, self.sign2,
                self.clamp_min_decay,
            )
            k_f = torch.fft.rfft(k.float(), n=n_fft)
        else:
            if (
                self._cached_k_f is None
                or self._cached_n_fft != n_fft
                or self._cached_k_f.device != x.device
            ):
                t = self._get_time(seqlen, x.device)
                k = _build_osc_kernel(
                    t,
                    self.log_alpha, self.log_omega, self.phi, self.log_amp,
                    self.log_beta, self.log_amp2, self.sign2,
                    self.clamp_min_decay,
                )
                self._cached_k_f = torch.fft.rfft(k.float(), n=n_fft)
                self._cached_n_fft = n_fft
            k_f = self._cached_k_f

        y = torch.fft.irfft(
            torch.fft.rfft(x_ct, n=n_fft) * k_f.unsqueeze(0),
            n=n_fft,
        )[:, :, :seqlen]

        return y.transpose(1, 2).to(dtype=x_dtype)

# =============================================================================
# Mixer Block
# =============================================================================


class HarmonicMixerBlock(nn.Module):
    """Gated mixer that uses the oscillator filter for sequence mixing."""

    def __init__(self, dim: int, config: BCLM1ModelConfig) -> None:
        super().__init__()
        self.norm = RMSNorm(dim)
        self.in_proj = nn.Linear(dim, 2 * dim, bias=False)
        self.gate_bias = nn.Parameter(torch.zeros(dim))
        self.osc = CausalOscillatorFilter(
            dim=dim,
            n_pairs=config.osc_n_pairs,
            n_real=config.osc_n_real,
            clamp_min_decay=config.osc_clamp_min_decay,
            max_len=config.max_seq_len,
        )
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = self.norm(x)
        u, z = self.in_proj(x).chunk(2, dim=-1)
        u = self.osc(u)
        out = u * (z + self.gate_bias)
        out = self.dropout(out)
        return self.norm(residual + out)


# =============================================================================
# Feed-Forward Network
# =============================================================================


class GatedFeedForward(nn.Module):
    """SwiGLU-style gated FFN in [B, T, C] layout."""

    def __init__(self, dim: int, dropout: float) -> None:
        super().__init__()
        self.norm = RMSNorm(dim)

        hidden_dim = ((16 * dim) // 3 + 63) // 64 * 64
        _check_tensor_core_alignment(hidden_dim, "ffn_hidden_dim")

        self.proj_in = nn.Linear(dim, 2 * hidden_dim, bias=False)
        self.proj_out = nn.Linear(hidden_dim, dim, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = self.norm(x)
        v, g = self.proj_in(x).chunk(2, dim=-1)
        x = self.proj_out(F.gelu(g, approximate="tanh") * v)
        return residual + self.dropout(x)


# =============================================================================
# Attention
# =============================================================================


class _SelfAttentionBase(nn.Module):
    """Shared projections, norm, and reshape logic for GQA attention."""

    def __init__(
        self,
        dim: int,
        n_heads: int,
        kv_heads: int,
        dropout: float,
    ) -> None:
        super().__init__()
        if dim % n_heads != 0:
            raise ValueError("embed_dim must be divisible by attn_heads")
        if kv_heads <= 0 or kv_heads > n_heads:
            raise ValueError("kv_heads must be in [1, attn_heads]")
        if n_heads % kv_heads != 0:
            raise ValueError("attn_heads must be divisible by kv_heads")

        self.dim = dim
        self.n_heads = n_heads
        self.kv_heads = kv_heads
        self.head_dim = dim // n_heads
        self.kv_dim = kv_heads * self.head_dim

        _check_tensor_core_alignment(self.head_dim, "head_dim")

        self.norm = RMSNorm(dim)
        self.qkv = nn.Linear(dim, dim + 2 * self.kv_dim, bias=False)
        self.out_proj = nn.Linear(dim, dim, bias=False)
        self.dropout = nn.Dropout(dropout)

    def _attention(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        p_dropout: float,
    ) -> torch.Tensor:
        raise NotImplementedError

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = self.norm(x)

        qkv = self.qkv(x)
        q, k, v = qkv.split([self.dim, self.kv_dim, self.kv_dim], dim=-1)

        bsz, seqlen, _ = q.shape
        q = q.view(bsz, seqlen, self.n_heads, self.head_dim)
        k = k.view(bsz, seqlen, self.kv_heads, self.head_dim)
        v = v.view(bsz, seqlen, self.kv_heads, self.head_dim)

        p_drop = self.dropout.p if self.training else 0.0
        out = self._attention(q, k, v, p_drop)

        out = out.reshape(bsz, seqlen, self.dim)
        out = self.out_proj(out)
        return residual + self.dropout(out)


def _sdpa_gqa(
    qh: torch.Tensor,
    kh: torch.Tensor,
    vh: torch.Tensor,
    p_dropout: float,
    attn_mask: Optional[torch.Tensor] = None,
    is_causal: bool = False,
) -> torch.Tensor:
    if kh.size(1) != qh.size(1):
        if _sdpa_supports_enable_gqa():
            return F.scaled_dot_product_attention(
                qh, kh, vh,
                attn_mask=attn_mask,
                dropout_p=p_dropout,
                is_causal=is_causal,
                enable_gqa=True,
            )
        kh = _repeat_kv_heads(kh, qh.size(1))
        vh = _repeat_kv_heads(vh, qh.size(1))
    return F.scaled_dot_product_attention(
        qh, kh, vh,
        attn_mask=attn_mask,
        dropout_p=p_dropout,
        is_causal=is_causal,
    )


class SlidingWindowSelfAttention(_SelfAttentionBase):
    """
    Local sliding-window causal attention with GQA.

    Uses xFormers when available, falls back to SDPA with a static mask.
    """

    def __init__(
        self,
        dim: int,
        n_heads: int,
        kv_heads: int,
        window_size: int,
        dropout: float,
        max_seq_len: int = 0,
    ) -> None:
        super().__init__(dim=dim, n_heads=n_heads, kv_heads=kv_heads, dropout=dropout)
        self.window_size = int(window_size)

        self._xf_bias = None
        self._xf_failed = False
        if _XFORMERS_AVAILABLE:
            try:
                self._xf_bias = (
                    fmha.attn_bias.LowerTriangularFromBottomRightMask()
                    .make_local_attention(self.window_size)
                )
            except Exception:
                self._xf_bias = None
                self._xf_failed = True

        if max_seq_len > 0:
            self.register_buffer(
                "_attn_mask",
                self._make_mask(max_seq_len, self.window_size),
                persistent=False,
            )
        else:
            self.register_buffer(
                "_attn_mask",
                torch.empty(0, dtype=torch.bool),
                persistent=False,
            )

    @staticmethod
    def _make_mask(seqlen: int, window_size: int) -> torch.Tensor:
        i = torch.arange(seqlen).unsqueeze(1)
        j = torch.arange(seqlen).unsqueeze(0)
        wl = max(window_size - 1, 0)
        return (j <= i) & (j >= (i - wl))

    def _get_mask(self, seqlen: int, device: torch.device) -> torch.Tensor:
        if self._attn_mask.numel() > 0 and seqlen <= self._attn_mask.size(0):
            return self._attn_mask[:seqlen, :seqlen]
        return self._make_mask(seqlen, self.window_size).to(device)

    def _attention(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        p_dropout: float,
    ) -> torch.Tensor:
        if (
            not self._xf_failed
            and self._xf_bias is not None
            and _can_use_xformers_attention(q, k, v)
        ):
            try:
                k_xf, v_xf = _expand_kv_for_xformers(q, k, v)
                return xops.memory_efficient_attention(
                    q, k_xf, v_xf, attn_bias=self._xf_bias, p=p_dropout,
                )
            except Exception:
                self._xf_failed = True

        seqlen = q.size(1)
        mask = self._get_mask(seqlen, q.device)
        out = _sdpa_gqa(
            q.transpose(1, 2),
            k.transpose(1, 2),
            v.transpose(1, 2),
            p_dropout,
            attn_mask=mask,
        )
        return out.transpose(1, 2).contiguous()


class GlobalSelfAttention(_SelfAttentionBase):
    """
    Full (global) causal attention with GQA.

    Uses xFormers when available, falls back to SDPA with is_causal=True.
    """

    def __init__(
        self,
        dim: int,
        n_heads: int,
        kv_heads: int,
        dropout: float,
    ) -> None:
        super().__init__(dim=dim, n_heads=n_heads, kv_heads=kv_heads, dropout=dropout)

        self._xf_bias = None
        self._xf_failed = False
        if _XFORMERS_AVAILABLE:
            try:
                self._xf_bias = fmha.attn_bias.LowerTriangularMask()
            except Exception:
                self._xf_bias = None
                self._xf_failed = True

    def _attention(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        p_dropout: float,
    ) -> torch.Tensor:
        if (
            not self._xf_failed
            and self._xf_bias is not None
            and k.size(-2) == q.size(-2)
            and _can_use_xformers_attention(q, k, v)
        ):
            try:
                return xops.memory_efficient_attention(
                    q, k, v, attn_bias=self._xf_bias, p=p_dropout,
                )
            except Exception:
                self._xf_failed = True

        out = _sdpa_gqa(
            q.transpose(1, 2),
            k.transpose(1, 2),
            v.transpose(1, 2),
            p_dropout,
            is_causal=True,
        )
        return out.transpose(1, 2).contiguous()


# =============================================================================
# Transformer Layer
# =============================================================================


class BCLM1Layer(nn.Module):
    """Single BCLM-1 layer: Mixer → (optional) Attention → FFN."""

    def __init__(
        self,
        config: BCLM1ModelConfig,
        layer_idx: int,
        local_attn_layer_set: Set[int],
        global_attn_layer_set: Set[int],
    ) -> None:
        super().__init__()
        dim = config.embed_dim
        n_heads = int(config.attn_heads)
        kv_heads = int(getattr(config, "attn_kv_heads", n_heads))

        self.mixer = HarmonicMixerBlock(dim, config)

        if layer_idx in local_attn_layer_set:
            self.attn: nn.Module = SlidingWindowSelfAttention(
                dim=dim,
                n_heads=n_heads,
                kv_heads=kv_heads,
                window_size=config.attn_window_size,
                dropout=config.dropout,
                max_seq_len=config.max_seq_len,
            )
        elif layer_idx in global_attn_layer_set:
            self.attn = GlobalSelfAttention(
                dim=dim,
                n_heads=n_heads,
                kv_heads=kv_heads,
                dropout=config.dropout,
            )
        else:
            self.attn = nn.Identity()

        self.ffn = GatedFeedForward(dim, config.dropout)
        self._use_checkpoint: bool = False

    def _forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.mixer(x)
        x = self.attn(x)
        x = self.ffn(x)
        return x

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._use_checkpoint and self.training:
            return grad_checkpoint(self._forward, x, use_reentrant=False)
        return self._forward(x)


# =============================================================================
# Main Model
# =============================================================================


class BCLM1Model(nn.Module):
    """
    BCLM-1 language model.

    Components:
        - Oscillator-based sequence mixing (HarmonicMixerBlock)
        - Sparse sliding-window / global causal attention (GQA)
        - Gated feed-forward networks (SwiGLU-style)
        - Bigram hash embedding injection
    """

    def __init__(self, cfg: Union[BCLM1ModelConfig, dict]) -> None:
        super().__init__()
        config = (
            cfg
            if isinstance(cfg, BCLM1ModelConfig)
            else BCLM1ModelConfig.from_dict(cfg)
        )
        self.config = config
        dim = config.embed_dim

        _check_tensor_core_alignment(dim, "embed_dim")
        _check_tensor_core_alignment(config.vocab_size, "vocab_size")

        self.tok_emb = nn.Embedding(config.vocab_size, dim)
        self.input_mixer = CausalConv(
            dim,
            kernel_size=config.conv_kernel_size,
            dropout=config.dropout,
        )

        self._bigram_vocab_size = config.vocab_size * config.bigram_table_factor
        self.bigram_embed = nn.Embedding(self._bigram_vocab_size, dim)
        nn.init.zeros_(self.bigram_embed.weight)

        self.bigram_lambdas = nn.Parameter(torch.full((config.n_layers,), 0.1))

        local_attn_layers = set(config.local_attn_layers)
        global_attn_layers = set(config.global_attn_layers)
        overlap = local_attn_layers & global_attn_layers
        if overlap:
            raise ValueError(
                f"Layers cannot be both local and global attention: {sorted(overlap)}"
            )

        self.layers = nn.ModuleList(
            [
                BCLM1Layer(config, i, local_attn_layers, global_attn_layers)
                for i in range(config.n_layers)
            ]
        )

        self.norm_f = RMSNorm(dim)
        self.head = nn.Linear(dim, config.vocab_size, bias=False)

        self._ce_chunk_size: int = 0

    @staticmethod
    def _compute_bigram_hash(x: torch.Tensor, bigram_vocab_size: int) -> torch.Tensor:
        """
        Map each (prev_token, curr_token) pair to a compact hash index.

        Uses two large odd integers with XOR + modulo for a roughly
        uniform distribution at O(1) cost.
        """
        rand_int_1 = 36313
        rand_int_2 = 27191
        mod = bigram_vocab_size - 1

        result = torch.empty_like(x)
        result[:, 0] = mod
        result[:, 1:] = (
            torch.bitwise_xor(rand_int_1 * x[:, 1:], rand_int_2 * x[:, :-1]) % mod
        )
        return result

    def set_ce_chunk_size(self, n: int) -> None:
        """Enable chunked cross-entropy (0 = disabled, typical: 4096)."""
        self._ce_chunk_size = max(0, int(n))

    def enable_gradient_checkpointing(
        self,
        enable: bool = True,
        layer_indices: Optional[list[int]] = None,
    ) -> None:
        """Toggle gradient checkpointing per layer."""
        for i, layer in enumerate(self.layers):
            if layer_indices is None or i in layer_indices:
                layer._use_checkpoint = enable

    def _chunked_cross_entropy(
        self,
        hidden: torch.Tensor,
        targets: torch.Tensor,
    ) -> torch.Tensor:
        B, T, C = hidden.shape
        N = B * T
        chunk = self._ce_chunk_size

        h_flat = hidden.reshape(N, C)
        t_flat = targets.reshape(N)
        weight = self.head.weight

        if N <= chunk:
            return F.cross_entropy(
                F.linear(h_flat, weight), t_flat, ignore_index=-100,
            )

        loss_acc = h_flat.new_zeros((), dtype=torch.float32)
        valid_acc = h_flat.new_zeros((), dtype=torch.long)

        for start in range(0, N, chunk):
            end = min(start + chunk, N)
            logits_c = F.linear(h_flat[start:end], weight)
            t_c = t_flat[start:end]

            n_valid = (t_c != -100).sum()
            loss_acc = loss_acc + F.cross_entropy(
                logits_c, t_c, ignore_index=-100, reduction="sum",
            )
            valid_acc = valid_acc + n_valid

        return loss_acc / valid_acc.clamp_min(1).float()

    def forward(
        self,
        idx: torch.Tensor,
        targets: Optional[torch.Tensor] = None,
    ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
        """
        Args:
            idx:     [B, T] input token ids.
            targets: [B, T] target token ids (optional; -100 = ignore).

        Returns:
            (logits, loss) — logits is None when chunked CE is active
            and targets are provided.
        """
        seqlen = idx.size(1)
        if seqlen > self.config.max_seq_len:
            raise ValueError(
                f"seqlen={seqlen} exceeds max_seq_len={self.config.max_seq_len}"
            )

        x = self.tok_emb(idx)
        x = x + self.input_mixer(x)

        bigram_idx = self._compute_bigram_hash(idx, self._bigram_vocab_size)
        x0_bigram = self.bigram_embed(bigram_idx)

        for i, layer in enumerate(self.layers):
            x = x + self.bigram_lambdas[i] * x0_bigram
            x = layer(x)

        x = self.norm_f(x)

        if targets is not None and self._ce_chunk_size > 0:
            loss = self._chunked_cross_entropy(x, targets)
            return None, loss

        logits = self.head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.reshape(-1, logits.size(-1)),
                targets.reshape(-1),
                ignore_index=-100,
            )

        return logits, loss

    def num_parameters(self, exclude_embeddings: bool = False) -> int:
        """Total number of scalar parameters."""
        n = sum(p.numel() for p in self.parameters())
        if exclude_embeddings:
            n -= self.tok_emb.weight.numel()
        return n