Module 3. The Basic Transformer

from __future__ import annotations

import importlib.util
import subprocess
import sys

REQUIRED = {"torch": "torch", "matplotlib": "matplotlib"}
missing = [package for package, module in REQUIRED.items() if importlib.util.find_spec(module) is None]
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])

import hashlib
import json
import math
from dataclasses import dataclass
from numbers import Integral, Real
from typing import Literal

import matplotlib.pyplot as plt
import torch
from torch import Tensor, nn
import torch.nn.functional as F

SEED = 18
torch.manual_seed(SEED)
torch.set_num_threads(1)
torch.use_deterministic_algorithms(True)
print(f"torch={torch.__version__} · device=cpu · seed={SEED}")

import numpy as np
from matplotlib.colors import ListedColormap
BLUE = '#5E81AC'
GREEN = '#A3BE8C'
ORANGE = '#D08770'
INK = '#2E3440'
MUTED = '#4C566A'
BORDER = '#D8DEE9'

def plot_style():
    import matplotlib.pyplot as plt
    plt.rcParams.update({
        "font.family": "DejaVu Sans", "font.size": 12,
        "text.color": INK, "axes.labelcolor": INK, "axes.titlecolor": INK,
        "axes.titlesize": 17, "axes.titleweight": "bold", "axes.titlelocation": "left",
        "axes.labelsize": 13, "xtick.labelsize": 11, "ytick.labelsize": 11,
        "xtick.color": MUTED, "ytick.color": MUTED,
        "axes.edgecolor": BORDER, "axes.linewidth": .7,
        "axes.facecolor": "white", "figure.facecolor": "white", "savefig.facecolor": "white",
        "grid.color": BORDER, "grid.linewidth": .7, "axes.axisbelow": True,
        "legend.fontsize": 11, "legend.frameon": False, "savefig.dpi": 150,
        "xtick.major.size": 0, "ytick.major.size": 0,
    })


def positional_plot(pe,language):
    plot_style();ru=language=='ru'
    fig,axes=plt.subplots(2,1,figsize=(11,6.3),layout='constrained',sharex=True)
    for ax,dims,title in zip(axes,((0,1),(4,5)),(('Быстрые координаты','Медленные координаты') if ru else ('Fast coordinates','Slow coordinates'))):
        for dim,col,style in zip(dims,(BLUE,ORANGE),('-','--')):
            ax.plot(np.arange(len(pe)),pe[:,dim],color=col,ls=style,lw=2,label=('координата ' if ru else 'coordinate ')+str(dim))
        ax.set_title(title,fontsize=14);ax.set_ylim(-1.1,1.1);ax.set_xlim(0,len(pe)-1);ax.set_ylabel('PE');ax.grid(alpha=.6);ax.legend(ncol=2,loc='lower right',bbox_to_anchor=(1,1.02))
    axes[-1].set_xlabel('Позиция (с нуля)' if ru else 'Position (zero-based)')
    fig.suptitle('Синусоидальные позиции · ширина D=8' if ru else 'Sinusoidal positions · width D=8',fontsize=17,fontweight='bold')
    return fig


def mask_plot(language):
    plot_style();ru=language=='ru'
    padding=np.ones((5,5),dtype=bool);padding[:,3:]=False
    causal=np.tril(np.ones((5,5),dtype=bool))
    cross=np.ones((4,6),dtype=bool);cross[:,5]=False
    masks=(padding,causal,causal&padding,cross)
    titles=('PAD-ключи скрыты','Будущее скрыто','Оба ограничения','Cross-attention') if ru else ('Hide padding keys','Hide future keys','Combine both','Cross-attention')
    fig,axes=plt.subplots(2,2,figsize=(11,8.8),layout='constrained')
    for ax,m,title in zip(axes.flat,masks,titles):
        ax.imshow(m,cmap=ListedColormap(['#ECEFF4',GREEN]),vmin=0,vmax=1)
        ax.set_title(title,fontsize=15);ax.set_xlabel('Ключи' if ru else 'Keys');ax.set_ylabel('Запросы' if ru else 'Queries')
        ax.set_xticks(range(m.shape[1]));ax.set_yticks(range(m.shape[0]))
        for i,j in np.ndindex(m.shape):ax.text(j,i,'1' if m[i,j] else '0',ha='center',va='center',fontsize=12,color=INK)
    fig.suptitle('1 — разрешено; 0 — запрет до softmax' if ru else '1 = allowed; 0 = blocked before softmax',fontsize=17,fontweight='bold')
    return fig


def loss_plot(losses,language):
    plot_style();ru=language=='ru'
    fig,ax=plt.subplots(figsize=(11,4.8),layout='constrained')
    ax.plot(losses,color=BLUE);ax.set_yscale('log');ax.grid(alpha=.6)
    ax.set_title('Маленькая модель учит период' if ru else 'A tiny model learns a period')
    ax.set_xlabel('Шаг оптимизации' if ru else 'Optimization step')
    ax.set_ylabel('NLL, нат/токен · лог. шкала' if ru else 'NLL, nats/token · log scale')
    return fig


def sampling_plot(scenario, language):
    plot_style()
    ru = language == 'ru'
    fig, ax = plt.subplots(figsize=(11, 5.8), layout='constrained')
    policies = ('temperature_1', 'temperature_2', 'top_k_2', 'top_p_09')
    labels = ('T = 1', 'T = 2', 'T = 1, top-k = 2', 'T = 1, top-p = 0.9')
    colors = (BLUE, ORANGE, GREEN, '#8055A6')
    positions = np.arange(len(scenario['logits']))
    width = 0.19
    for index, (policy, label, color) in enumerate(zip(policies, labels, colors)):
        ax.bar(positions + (index - 1.5) * width, scenario[policy], width,
               label=label, color=color)
    ax.set_xticks(positions, [f'ID {i}' for i in positions])
    ax.set_ylim(0, 0.82)
    ax.set_ylabel('Вероятность после нормировки' if ru else 'Renormalized probability')
    ax.set_title('Одни логиты [3, 2, 1, 0] — четыре распределения' if ru else
                 'One logit vector [3, 2, 1, 0], four distributions', pad=62)
    ax.legend(ncol=2, loc='lower center', bbox_to_anchor=(0.5, 1.01))
    ax.grid(axis='y', alpha=0.6)
    ax.set_axisbelow(True)
    return fig

REFERENCE_JSON = '{"attention":{"allowed":[[true,false,false],[true,true,false],[true,true,true]],"k":[[1.0,0.0],[0.0,1.0],[1.0,1.0]],"output":[[1.0,0.0],[0.3302384506733431,1.3395230986533138],[1.7587246087113846,1.0]],"q":[[1.0,0.0],[0.0,1.0],[1.0,1.0]],"scores":[[0.7071067811865475,null,null],[0.0,0.7071067811865475,null],[0.7071067811865475,0.7071067811865475,1.414213562373095]],"v":[[1.0,0.0],[0.0,2.0],[3.0,1.0]],"weights":[[1.0,0.0,0.0],[0.3302384506733431,0.6697615493266569,0.0],[0.2482550782577231,0.2482550782577231,0.5034898434845538]]},"layer_norm":{"input":[1.0,2.0,3.0,4.0],"mean":2.5,"normalized":[-1.3416407864998738,-0.4472135954999579,0.4472135954999579,1.3416407864998738],"variance":1.25},"next_token_nll":{"logits":[1.2,0.3,-0.4,2.0],"nll":0.5439104980979277,"probabilities":[0.2608237210556189,0.10604301152185992,0.05265940109153875,0.5804738663309824],"target":3},"parameter_floor":{"d_ff":2048,"d_model":512,"decoder_layer":4194304,"encoder_layer":3145728,"ffn_share_encoder":0.6666666666666666,"ffn_weights":2097152,"mha_weights":1048576},"positional_encoding":{"d_model":4,"length":5,"table":[[0.0,1.0,0.0,1.0],[0.8414709848078965,0.5403023058681398,0.009999833334166664,0.9999500004166653],[0.9092974268256817,-0.4161468365471424,0.01999866669333308,0.9998000066665778],[0.1411200080598672,-0.9899924966004454,0.02999550020249566,0.9995500337489875],[-0.7568024953079282,-0.6536436208636119,0.03998933418663416,0.9992001066609779]]},"scaling_experiment":{"d_k":64,"relative_tolerance":0.08,"samples":50000,"scaled_variance_target":1.0,"seed":18,"unscaled_variance_target":64.0},"schema_version":1,"tiny_lm":{"accuracy_threshold":0.99,"context":12,"continuation_accuracy_threshold":0.95,"learning_rate":0.03,"loss_threshold":0.02,"new_tokens":15,"period":[0,1,2,3,4,5],"seed":18,"steps":121,"training_repeats":20,"vocab_size":6}}'
EXPECTED_SHA256 = 'e608c5bb68921c5a68e040a1e90f1bf1b9565a62a1b24435429707d4f4e31957'
actual_sha256 = hashlib.sha256(REFERENCE_JSON.encode("utf-8")).hexdigest()
assert actual_sha256 == EXPECTED_SHA256
REFERENCE = json.loads(REFERENCE_JSON)
assert REFERENCE["schema_version"] == 1
print('Embedded Module 3 contract verified', actual_sha256)

def _positive_int(name: str, value: object) -> int:
    """Return a strictly positive integer and reject booleans."""

    if isinstance(value, bool) or not isinstance(value, Integral):
        raise TypeError(f"{name} must be an integer")
    result = int(value)
    if result <= 0:
        raise ValueError(f"{name} must be positive")
    return result


def _non_negative_int(name: str, value: object) -> int:
    """Return a non-negative integer and reject booleans."""

    if isinstance(value, bool) or not isinstance(value, Integral):
        raise TypeError(f"{name} must be an integer")
    result = int(value)
    if result < 0:
        raise ValueError(f"{name} must be non-negative")
    return result


def expect_raises(error_type: type[BaseException], fn, *args, **kwargs) -> str:
    """Run ``fn`` and return the expected exception message."""

    try:
        fn(*args, **kwargs)
    except error_type as exc:
        return str(exc)
    raise AssertionError(f"expected {error_type.__name__}")


def _finite_positive_real(name: str, value: object) -> float:
    """Return a positive finite scalar."""

    if isinstance(value, bool) or not isinstance(value, Real):
        raise TypeError(f"{name} must be a real scalar")
    result = float(value)
    if not math.isfinite(result) or result <= 0.0:
        raise ValueError(f"{name} must be finite and positive")
    return result


def _floating_tensor(name: str, value: object, *, min_ndim: int) -> Tensor:
    """Validate a finite floating-point tensor without silently changing shape."""

    if not isinstance(value, Tensor):
        raise TypeError(f"{name} must be a torch.Tensor")
    if value.ndim < min_ndim:
        raise ValueError(f"{name} must have at least {min_ndim} dimensions")
    if any(size <= 0 for size in value.shape):
        raise ValueError(f"{name} may not contain an empty axis")
    if not value.is_floating_point():
        raise TypeError(f"{name} must use a floating-point dtype")
    if not torch.isfinite(value).all():
        raise ValueError(f"{name} must contain only finite values")
    return value

def sinusoidal_encoding(
    length: int,
    d_model: int,
    *,
    base: float = 10_000.0,
    device: torch.device | str | None = None,
    dtype: torch.dtype = torch.float64,
) -> Tensor:
    """Return the Vaswani sinusoidal table with shape ``[length, d_model]``.

    This paired implementation requires an even model width.  Position is the first
    axis; sine occupies even coordinates and cosine occupies odd coordinates.
    """

    length = _positive_int("length", length)
    d_model = _positive_int("d_model", d_model)
    base = _finite_positive_real("base", base)
    if base <= 1.0:
        raise ValueError("base must be greater than 1")
    if d_model % 2:
        raise ValueError("d_model must be even for paired sine/cosine coordinates")
    if dtype not in {torch.float32, torch.float64}:
        raise TypeError("dtype must be torch.float32 or torch.float64")

    position = torch.arange(length, device=device, dtype=dtype)[:, None]
    pair_index = torch.arange(0, d_model, 2, device=device, dtype=dtype)
    divisor = torch.pow(torch.tensor(base, device=device, dtype=dtype), pair_index / d_model)
    table = torch.empty((length, d_model), device=device, dtype=dtype)
    table[:, 0::2] = torch.sin(position / divisor)
    table[:, 1::2] = torch.cos(position / divisor)
    return table


def causal_allowed(
    t_q: int,
    t_k: int,
    *,
    query_offset: int = 0,
    device: torch.device | str | None = None,
) -> Tensor:
    """Return a boolean causal mask of shape ``[T_q, T_k]``.

    ``query_offset`` is the absolute position of local query row zero.  It is zero for
    ordinary self-attention and equals the cached-prefix length for incremental decode.
    ``True`` means that the query may read the key.
    """

    t_q = _positive_int("t_q", t_q)
    t_k = _positive_int("t_k", t_k)
    if isinstance(query_offset, bool) or not isinstance(query_offset, Integral):
        raise TypeError("query_offset must be an integer")
    query_offset = int(query_offset)
    if query_offset < 0:
        raise ValueError("query_offset must be non-negative")
    q_pos = query_offset + torch.arange(t_q, device=device)[:, None]
    k_pos = torch.arange(t_k, device=device)[None, :]
    return k_pos <= q_pos

pe = sinusoidal_encoding(81, 8)
assert pe.shape == (81, 8)
torch.testing.assert_close(pe[0, 0::2], torch.zeros(4, dtype=pe.dtype))
torch.testing.assert_close(pe[0, 1::2], torch.ones(4, dtype=pe.dtype))
positional_ref = REFERENCE["positional_encoding"]
positional_anchor = sinusoidal_encoding(
    positional_ref["length"],
    positional_ref["d_model"],
)
torch.testing.assert_close(
    positional_anchor,
    torch.tensor(positional_ref["table"], dtype=torch.float64),
    atol=1e-12,
    rtol=0,
)

fig = positional_plot(pe.numpy(), 'en')
plt.show()
plt.close(fig)
print('Sinusoidal table and position zero verified')

@dataclass(frozen=True)
class AttentionResult:
    """Transparent outputs of one scaled dot-product attention call."""

    output: Tensor
    weights: Tensor
    scores: Tensor
    empty_rows: Tensor


def scaled_dot_product_attention(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    allowed: Tensor | None = None,
    *,
    empty_row: EmptyRowPolicy = "error",
) -> AttentionResult:
    """Compute numerically stable educational scaled dot-product attention.

    Shapes:
        ``q [..., T_q, d_k]``
        ``k [..., T_k, d_k]``
        ``v [..., T_k, d_v]``
        ``allowed`` broadcasts to ``[..., T_q, T_k]`` and uses ``True = allowed``.

    The calculation uses at least float32 for logits.  A fully masked query row is not
    a probability distribution: ``empty_row='error'`` rejects it, while ``'zero'``
    returns zero weights and a zero output for explicitly padded query rows.
    """

    if empty_row not in {"error", "zero"}:
        raise ValueError("empty_row must be 'error' or 'zero'")
    q = _floating_tensor("q", q, min_ndim=2)
    k = _floating_tensor("k", k, min_ndim=2)
    v = _floating_tensor("v", v, min_ndim=2)
    if not (q.device == k.device == v.device):
        raise ValueError("q, k, and v must be on the same device")
    if q.shape[-1] != k.shape[-1]:
        raise ValueError("q and k must share d_k")
    if k.shape[-2] != v.shape[-2]:
        raise ValueError("k and v must share T_k")
    try:
        batch_shape = torch.broadcast_shapes(q.shape[:-2], k.shape[:-2], v.shape[:-2])
    except RuntimeError as exc:
        raise ValueError("q/k/v batch dimensions are not broadcast-compatible") from exc

    compute_dtype = torch.float64 if torch.float64 in {q.dtype, k.dtype, v.dtype} else torch.float32
    qf, kf, vf = (tensor.to(dtype=compute_dtype) for tensor in (q, k, v))
    # V can introduce a batch axis too; the mask belongs to the joint batch shape.
    qf, kf, vf = (t.expand(*batch_shape, *t.shape[-2:]) for t in (qf, kf, vf))
    scores = (qf @ kf.transpose(-2, -1)) / math.sqrt(q.shape[-1])
    if not torch.isfinite(scores).all():
        raise FloatingPointError("q @ k.T overflowed; rescale the inputs")

    if allowed is None:
        allowed_b = torch.ones_like(scores, dtype=torch.bool)
    else:
        if not isinstance(allowed, Tensor):
            raise TypeError("allowed must be a torch.Tensor")
        if allowed.dtype is not torch.bool:
            raise TypeError("allowed must use bool dtype")
        try:
            allowed_b = torch.broadcast_to(allowed.to(device=scores.device), scores.shape)
        except RuntimeError as exc:
            raise ValueError(
                f"allowed shape {tuple(allowed.shape)} does not broadcast to {tuple(scores.shape)}"
            ) from exc

    empty_rows = ~allowed_b.any(dim=-1)
    if empty_row == "error" and empty_rows.any():
        locations = torch.nonzero(empty_rows, as_tuple=False)[:8].cpu().tolist()
        raise ValueError(f"attention contains fully masked rows: {locations}")

    masked_scores = scores.masked_fill(~allowed_b, float("-inf"))
    row_has_support = ~empty_rows.unsqueeze(-1)
    row_max = torch.where(
        row_has_support,
        masked_scores.max(dim=-1, keepdim=True).values,
        torch.zeros_like(masked_scores[..., :1]),
    )
    exponentials = torch.where(
        allowed_b,
        torch.exp(masked_scores - row_max),
        torch.zeros_like(masked_scores),
    )
    denominator = exponentials.sum(dim=-1, keepdim=True)
    safe_denominator = torch.where(denominator > 0, denominator, torch.ones_like(denominator))
    weights = exponentials / safe_denominator
    output = weights @ vf
    if not torch.isfinite(weights).all() or not torch.isfinite(output).all():
        raise FloatingPointError("attention produced a non-finite weight or output")
    return AttentionResult(output=output, weights=weights, scores=masked_scores, empty_rows=empty_rows)

attention_ref = REFERENCE["attention"]
Q = torch.tensor(attention_ref["q"], dtype=torch.float64)
K = torch.tensor(attention_ref["k"], dtype=torch.float64)
V = torch.tensor(attention_ref["v"], dtype=torch.float64)
allowed = torch.tensor(attention_ref["allowed"], dtype=torch.bool)
result = scaled_dot_product_attention(Q, K, V, allowed)
expected_weights = torch.tensor(attention_ref["weights"], dtype=torch.float64)
expected_output = torch.tensor(attention_ref["output"], dtype=torch.float64)
torch.testing.assert_close(result.weights, expected_weights, atol=1e-12, rtol=0)
torch.testing.assert_close(result.output, expected_output, atol=1e-12, rtol=0)
torch.testing.assert_close(result.weights.sum(-1), torch.ones(3, dtype=torch.float64))
assert torch.all(result.weights[~allowed] == 0)

empty_mask = allowed.clone()
empty_mask[1, :] = False
try:
    scaled_dot_product_attention(Q, K, V, empty_mask)
    raise AssertionError("default policy must reject a fully masked row")
except ValueError as exc:
    assert "fully masked" in str(exc)
zero_result = scaled_dot_product_attention(Q, K, V, empty_mask, empty_row="zero")
assert torch.all(zero_result.weights[1] == 0)
assert torch.all(zero_result.output[1] == 0)
assert not torch.isnan(zero_result.weights).any()
print("weights:\n", result.weights)
print("head output Z:\n", result.output)
print('Numerical attention and both empty-row policies verified')

def check_attention_contract(attention):
    """Compare outputs/gradients, batch broadcasting, and cached-query geometry.

    The callable returns an object with .output and .weights. No repository imports
    are needed here, so the notebook builder can embed this complete function.
    """
    import torch
    import torch.nn.functional as F
    torch.manual_seed(303)
    q=torch.randn(2,2,3,4,dtype=torch.float64,requires_grad=True)
    k=torch.randn(2,2,5,4,dtype=torch.float64,requires_grad=True)
    v=torch.randn(2,2,5,6,dtype=torch.float64,requires_grad=True)
    allowed=torch.arange(5)[None,:] <= (2+torch.arange(3))[:,None]
    explicit=attention(q,k,v,allowed).output
    library=F.scaled_dot_product_attention(q,k,v,attn_mask=allowed,dropout_p=0.0)
    torch.testing.assert_close(explicit,library,atol=1e-12,rtol=1e-12)
    probe=torch.randn_like(explicit)
    explicit_grad=torch.autograd.grad((explicit*probe).sum(),(q,k,v))
    library_grad=torch.autograd.grad((library*probe).sum(),(q,k,v))
    for a,b in zip(explicit_grad,library_grad):
        torch.testing.assert_close(a,b,atol=1e-11,rtol=1e-11)

    # A mask can vary across a batch axis introduced only by V.
    bq=torch.zeros(1,2,1,dtype=torch.float64)
    bk=torch.zeros(1,3,1,dtype=torch.float64)
    bv=torch.tensor([[[1.],[2.],[3.]],[[4.],[5.],[6.]]],dtype=torch.float64)
    bm=torch.tensor([[[1,0,0],[0,1,0]],[[0,0,1],[1,1,0]]],dtype=torch.bool)
    batched=attention(bq,bk,bv,bm)
    torch.testing.assert_close(batched.output,torch.tensor([[[1.],[2.]],[[6.],[4.5]]],dtype=torch.float64))
    assert batched.weights.shape==(2,2,3)

    # Explicit zero rows also have finite gradients, rather than hidden NaNs.
    zq=q.detach().clone().requires_grad_()
    zk=k.detach().clone().requires_grad_()
    zv=v.detach().clone().requires_grad_()
    zm=allowed.clone();zm[1]=False
    zero=attention(zq,zk,zv,zm,empty_row='zero')
    assert torch.all(zero.output[...,1,:]==0)
    zg=torch.autograd.grad(zero.output.square().sum(),(zq,zk,zv))
    assert all(torch.isfinite(g).all() for g in zg)
    assert torch.all(zg[0][...,1,:]==0)

    # Last two queries have absolute positions 3 and 4, not local positions 0 and 1.
    fq=torch.randn(5,4,dtype=torch.float64)
    fk=torch.randn(5,4,dtype=torch.float64)
    fv=torch.randn(5,2,dtype=torch.float64)
    cm=torch.ones(5,5,dtype=torch.bool).tril()
    full=attention(fq,fk,fv,cm).output
    incremental=attention(fq[3:],fk,fv,cm[3:]).output
    torch.testing.assert_close(incremental,full[3:],atol=1e-12,rtol=1e-12)
    wrong=attention(fq[3:],fk,fv,torch.ones(2,5,dtype=torch.bool).tril()).output
    assert float((wrong-full[3:]).abs().max()) > 1e-3

    # Permuting unmasked input permutes output; it does not leave output rows fixed.
    perm=torch.tensor([2,0,4,1,3])
    original=attention(fq,fk,fv).output
    permuted=attention(fq[perm],fk[perm],fv[perm]).output
    torch.testing.assert_close(permuted,original[perm],atol=1e-12,rtol=1e-12)
    return {'output_max_error':float((explicit-library).abs().max().detach()),
            'gradient_max_error':max(float((a-b).abs().max()) for a,b in zip(explicit_grad,library_grad)),
            'wrong_cache_mask_max_error':float((wrong-full[3:]).abs().max()),
            'batch_broadcast':'PASS','zero_row_gradients':'PASS','permutation':'PASS'}

print(check_attention_contract(scaled_dot_product_attention))

scaling = REFERENCE["scaling_experiment"]
generator = torch.Generator().manual_seed(scaling["seed"])
q_mc = torch.randn((scaling["samples"], scaling["d_k"]), generator=generator)
k_mc = torch.randn((scaling["samples"], scaling["d_k"]), generator=generator)
dots = (q_mc * k_mc).sum(dim=-1)
unscaled_variance = float(dots.var(unbiased=False))
scaled_variance = float((dots / math.sqrt(scaling["d_k"])).var(unbiased=False))
tolerance = scaling["relative_tolerance"]
assert abs(unscaled_variance / scaling["unscaled_variance_target"] - 1.0) < tolerance
assert abs(scaled_variance / scaling["scaled_variance_target"] - 1.0) < tolerance
del q_mc, k_mc, dots

ordinary = causal_allowed(5, 5)
local_decode_wrong = causal_allowed(1, 5, query_offset=0)
cached_decode = causal_allowed(1, 5, query_offset=4)
assert ordinary.sum() == 15
assert local_decode_wrong.sum() == 1
assert cached_decode.sum() == 5

padding = torch.ones((5, 5), dtype=torch.bool)
padding[:, 3:] = False
combined = ordinary & padding
fig = mask_plot('en')
plt.show()
plt.close(fig)
print({"unscaled_var": unscaled_variance, "scaled_var": scaled_variance})
print('Scaling and cached-decode causal geometry verified')

class MultiHeadSelfAttention(nn.Module):
    """Minimal self-attention with explicit head geometry and padding semantics."""

    def __init__(self, d_model: int, n_heads: int, *, causal: bool = True) -> None:
        """Validate head geometry and create Q/K/V/O projections."""

        super().__init__()
        d_model = _positive_int("d_model", d_model)
        n_heads = _positive_int("n_heads", n_heads)
        if d_model % n_heads:
            raise ValueError("d_model must be divisible by n_heads")
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.causal = bool(causal)
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out = nn.Linear(d_model, d_model, bias=False)

    def forward(
        self,
        x: Tensor,
        *,
        token_valid: Tensor | None = None,
        empty_row: EmptyRowPolicy = "error",
        return_weights: bool = False,
    ) -> Tensor | tuple[Tensor, Tensor]:
        """Apply self-attention to ``x [B,T,D]``.

        ``token_valid [B,T]`` marks real tokens.  Invalid query rows are fully masked;
        use ``empty_row='zero'`` when padding rows should yield zero attention output.
        """

        x = _floating_tensor("x", x, min_ndim=3)
        if x.ndim != 3 or x.shape[-1] != self.d_model:
            raise ValueError(f"x must have shape [B,T,{self.d_model}]")
        batch, length, width = x.shape
        if token_valid is None:
            token_valid = torch.ones((batch, length), dtype=torch.bool, device=x.device)
        else:
            if not isinstance(token_valid, Tensor) or token_valid.dtype is not torch.bool:
                raise TypeError("token_valid must be a bool tensor")
            if token_valid.shape != (batch, length):
                raise ValueError(f"token_valid must have shape {(batch, length)}")
            token_valid = token_valid.to(device=x.device)

        q, k, v = self.qkv(x).chunk(3, dim=-1)

        def split_heads(tensor: Tensor) -> Tensor:
            """Reshape ``[B,T,D]`` into ``[B,H,T,d_head]`` without copying semantics."""

            return tensor.view(batch, length, self.n_heads, self.d_head).transpose(1, 2)

        q, k, v = map(split_heads, (q, k, v))
        allowed = token_valid[:, None, :, None] & token_valid[:, None, None, :]
        if self.causal:
            allowed &= causal_allowed(length, length, device=x.device)[None, None, :, :]
        result = scaled_dot_product_attention(q, k, v, allowed, empty_row=empty_row)
        merged = result.output.transpose(1, 2).contiguous().view(batch, length, width)
        output = self.out(merged.to(dtype=x.dtype))
        return (output, result.weights) if return_weights else output


class FeedForward(nn.Module):
    """Position-wise two-layer GELU feed-forward network."""

    def __init__(self, d_model: int, d_ff: int) -> None:
        """Create the two affine maps and intermediate GELU activation."""

        super().__init__()
        d_model = _positive_int("d_model", d_model)
        d_ff = _positive_int("d_ff", d_ff)
        self.net = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
        )

    def forward(self, x: Tensor) -> Tensor:
        """Preserve all leading axes and the residual width ``d_model``."""

        return self.net(x)


class DecoderBlock(nn.Module):
    """Minimal pre-norm decoder block with causal self-attention and an FFN."""

    def __init__(self, d_model: int, n_heads: int, d_ff: int) -> None:
        """Create two pre-norm residual branches with fixed model width."""

        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = MultiHeadSelfAttention(d_model, n_heads, causal=True)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = FeedForward(d_model, d_ff)

    def forward(self, x: Tensor) -> Tensor:
        """Return ``x + Attention(LN(x)) + FFN(LN(...))`` with the same shape."""

        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

torch.manual_seed(SEED)
mha = MultiHeadSelfAttention(d_model=32, n_heads=4)
x = torch.randn(2, 5, 32)
token_valid = torch.tensor([
    [True, True, True, True, True],
    [True, True, True, False, False],
])
y, weights = mha(x, token_valid=token_valid, empty_row="zero", return_weights=True)
assert y.shape == (2, 5, 32)
assert weights.shape == (2, 4, 5, 5)
assert torch.all(weights[1, :, 3:, :] == 0)
assert torch.all(y[1, 3:, :] == 0)
valid_query = token_valid[:, None, :, None].expand_as(weights[..., :1])
row_sums = weights.sum(-1, keepdim=True)
torch.testing.assert_close(row_sums[valid_query], torch.ones_like(row_sums[valid_query]), atol=1e-6, rtol=1e-6)
try:
    mha(x, token_valid=token_valid, empty_row="error")
    raise AssertionError("padding queries must trigger the default empty-row policy")
except ValueError as exc:
    assert "fully masked" in str(exc)
print("x", tuple(x.shape), "· weights", tuple(weights.shape), "· y", tuple(y.shape))
print('Multi-head reshape, padding, and normalization verified')

@dataclass(frozen=True)
class ParameterFloor:
    """Bias/norm-free parameter floor for the original Transformer sublayers."""

    mha_weights: int
    ffn_weights: int
    encoder_layer: int
    decoder_layer: int
    ffn_share_encoder: float


def transformer_parameter_floor(d_model: int, d_ff: int) -> ParameterFloor:
    """Count Q/K/V/O and two- or three-sublayer weights, excluding biases and norms."""

    d_model = _positive_int("d_model", d_model)
    d_ff = _positive_int("d_ff", d_ff)
    mha = 4 * d_model * d_model
    ffn = 2 * d_model * d_ff
    return ParameterFloor(
        mha_weights=mha,
        ffn_weights=ffn,
        encoder_layer=mha + ffn,
        decoder_layer=2 * mha + ffn,
        ffn_share_encoder=ffn / (mha + ffn),
    )

layer_ref = REFERENCE["layer_norm"]
x_ln = torch.tensor(layer_ref["input"], dtype=torch.float64)
mean = x_ln.mean()
variance = ((x_ln - mean) ** 2).mean()
manual = (x_ln - mean) / torch.sqrt(variance)
expected = torch.tensor(layer_ref["normalized"], dtype=torch.float64)
torch.testing.assert_close(manual, expected, atol=1e-12, rtol=0)
torch.testing.assert_close(F.layer_norm(x_ln, (4,), eps=0.0), expected, atol=1e-12, rtol=0)

block = DecoderBlock(32, 4, 64)
block_input = torch.randn(2, 5, 32)
block_output = block(block_input)
assert block_output.shape == block_input.shape
assert torch.isfinite(block_output).all()
print("mean", float(mean), "· variance", float(variance))
print('LayerNorm and decoder block verified')

class CausalityProbe(nn.Module):
    """Deterministic model whose unmasked output must depend on the future."""

    def __init__(self, *, causal: bool) -> None:
        """Create fixed embeddings and value-only attention projections."""

        super().__init__()
        width = 8
        self.embedding = nn.Embedding(11, width)
        self.attention = MultiHeadSelfAttention(width, 2, causal=causal)
        with torch.no_grad():
            values = torch.arange(11 * width, dtype=torch.float32).view(11, width) / 10.0
            self.embedding.weight.copy_(values)
            self.attention.qkv.weight.zero_()
            self.attention.qkv.weight[2 * width :, :].copy_(torch.eye(width))
            self.attention.out.weight.copy_(torch.eye(width))

    def forward(self, token_ids: Tensor) -> Tensor:
        """Return hidden states whose only cross-token path is attention."""

        hidden = self.embedding(token_ids)
        return hidden + self.attention(hidden)

base = torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.long)
changed_future = torch.tensor([[1, 2, 3, 4, 9]], dtype=torch.long)
causal_probe = CausalityProbe(causal=True).eval()
open_probe = CausalityProbe(causal=False).eval()
with torch.no_grad():
    causal_a, causal_b = causal_probe(base), causal_probe(changed_future)
    open_a, open_b = open_probe(base), open_probe(changed_future)
causal_early_diff = float((causal_a[:, :4] - causal_b[:, :4]).abs().max())
open_early_diff = float((open_a[:, :4] - open_b[:, :4]).abs().max())
assert causal_early_diff <= 1e-7
assert open_early_diff > 0.1
print({"causal_early_max_diff": causal_early_diff, "unmasked_early_max_diff": open_early_diff})
print('Deterministic future-leakage test passed')

class TinyLM(nn.Module):
    """One-block causal language model used only for a bounded overfit smoke test."""

    def __init__(
        self,
        *,
        vocab_size: int = 6,
        d_model: int = 24,
        n_heads: int = 4,
        d_ff: int = 48,
        max_length: int = 32,
    ) -> None:
        """Create a bounded one-block causal LM with tied token/output weights."""

        super().__init__()
        self.vocab_size = _positive_int("vocab_size", vocab_size)
        self.max_length = _positive_int("max_length", max_length)
        self.token_embedding = nn.Embedding(self.vocab_size, d_model)
        self.position_embedding = nn.Embedding(self.max_length, d_model)
        self.block = DecoderBlock(d_model, n_heads, d_ff)
        self.final_norm = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, self.vocab_size, bias=False)
        self.head.weight = self.token_embedding.weight

    def forward(self, token_ids: Tensor) -> Tensor:
        """Map integer IDs ``[B,T]`` to next-token logits ``[B,T,V]``."""

        if not isinstance(token_ids, Tensor) or token_ids.dtype is not torch.long:
            raise TypeError("token_ids must be a torch.long tensor")
        if token_ids.ndim != 2 or any(size <= 0 for size in token_ids.shape):
            raise ValueError("token_ids must have non-empty shape [B,T]")
        if token_ids.shape[1] > self.max_length:
            raise ValueError("sequence exceeds max_length")
        if int(token_ids.min()) < 0 or int(token_ids.max()) >= self.vocab_size:
            raise ValueError("token_ids contain an ID outside the vocabulary")
        positions = torch.arange(token_ids.shape[1], device=token_ids.device)
        hidden = self.token_embedding(token_ids) + self.position_embedding(positions)
        hidden = self.block(hidden)
        return self.head(self.final_norm(hidden))

    @torch.no_grad()
    def generate(self, token_ids: Tensor, new_tokens: int) -> Tensor:
        """Greedily append ``new_tokens`` IDs while cropping the model context."""

        new_tokens = _non_negative_int("new_tokens", new_tokens)
        generated = token_ids.clone()
        was_training = self.training
        self.eval()
        try:
            # Cropping resets learned positions; this is not KV-cached decoding.
            for _ in range(new_tokens):
                context = generated[:, -self.max_length :]
                next_id = self(context)[:, -1].argmax(dim=-1, keepdim=True)
                generated = torch.cat((generated, next_id), dim=1)
        finally:
            self.train(was_training)
        return generated

config = REFERENCE["tiny_lm"]
torch.manual_seed(config["seed"])
model = TinyLM(vocab_size=config["vocab_size"])
optimizer = torch.optim.AdamW(model.parameters(), lr=config["learning_rate"], weight_decay=0.0)
period = torch.tensor(config["period"] * config["training_repeats"], dtype=torch.long)
context = config["context"]
x_train = torch.stack([period[index : index + context] for index in range(len(period) - context)])
y_train = torch.stack([period[index + 1 : index + context + 1] for index in range(len(period) - context)])

losses = []
for _ in range(config["steps"]):
    logits = model(x_train)
    loss = F.cross_entropy(logits.reshape(-1, model.vocab_size), y_train.reshape(-1))
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    losses.append(float(loss.detach()))

with torch.no_grad():
    train_predictions = model(x_train).argmax(dim=-1)
    train_accuracy = float((train_predictions == y_train).float().mean())
assert losses[-1] < config["loss_threshold"]
assert train_accuracy >= config["accuracy_threshold"]
assert losses[-1] < losses[0] / 50

start = torch.tensor([[0, 1, 2]], dtype=torch.long)
generated = model.generate(start, config["new_tokens"])[0]
expected = torch.tensor([config["period"][i % len(config["period"])] for i in range(generated.numel())], dtype=torch.long)
prefix_length = start.shape[1]
continuation_accuracy = float((generated[prefix_length:] == expected[prefix_length:]).float().mean())
assert continuation_accuracy >= config["continuation_accuracy_threshold"]

fig = loss_plot(losses, 'en')
plt.show()
plt.close(fig)
print({"loss_start": losses[0], "loss_end": losses[-1], "train_accuracy": train_accuracy})
print("generated", generated)
print('Tiny LM overfit and generated the period')

nll_ref = REFERENCE["next_token_nll"]
logits = torch.tensor(nll_ref["logits"], dtype=torch.float64)
probabilities = torch.softmax(logits, dim=-1)
nll = -torch.log(probabilities[nll_ref["target"]])
torch.testing.assert_close(
    probabilities,
    torch.tensor(nll_ref["probabilities"], dtype=torch.float64),
    atol=1e-12,
    rtol=0,
)
assert abs(float(nll) - nll_ref["nll"]) < 1e-12

floor_ref = REFERENCE["parameter_floor"]
floor = transformer_parameter_floor(floor_ref["d_model"], floor_ref["d_ff"])
assert floor.mha_weights == floor_ref["mha_weights"]
assert floor.ffn_weights == floor_ref["ffn_weights"]
assert floor.encoder_layer == floor_ref["encoder_layer"]
assert floor.decoder_layer == floor_ref["decoder_layer"]
assert floor.ffn_share_encoder == floor_ref["ffn_share_encoder"]

def expect_error(error_type, function, *args, **kwargs):
    """Assert that one documented invalid input fails explicitly."""
    try:
        function(*args, **kwargs)
        raise AssertionError(f"expected {error_type.__name__}")
    except error_type:
        pass

expect_error(ValueError, sinusoidal_encoding, 5, 7)
expect_error(ValueError, causal_allowed, 1, 5, query_offset=-1)
expect_error(ValueError, scaled_dot_product_attention, Q[:, :1], K, V)
expect_error(TypeError, scaled_dot_product_attention, Q, K, V, torch.ones(3, 3))
expect_error(ValueError, MultiHeadSelfAttention, 30, 8)
expect_error(ValueError, model, torch.tensor([[0, 99]], dtype=torch.long))

print("probabilities", probabilities.tolist(), "· NLL", float(nll))
print(floor)
print('NLL, parameter floor, and negative tests passed')

"""A transparent scalar reference for next-token selection policies."""
import math
from numbers import Integral, Real


def token_distribution(logits, temperature=1.0, top_k=None, top_p=None):
    """Return probabilities for one finite logit vector; apply temperature, k, then p."""
    values = list(logits)
    if not values or any(isinstance(x, bool) or not isinstance(x, Real)
                         or not math.isfinite(x) for x in values):
        raise ValueError('logits must be a nonempty vector of finite real numbers')
    if isinstance(temperature, bool) or not isinstance(temperature, Real) or not math.isfinite(temperature) or temperature <= 0:
        raise ValueError('temperature must be finite and positive')
    if top_k is not None and (isinstance(top_k, bool) or not isinstance(top_k, Integral) or not 1 <= top_k <= len(values)):
        raise ValueError('top_k must be an integer in [1, vocabulary size]')
    if top_p is not None and (isinstance(top_p, bool) or not isinstance(top_p, Real) or not math.isfinite(top_p) or not 0 < top_p <= 1):
        raise ValueError('top_p must lie in (0, 1]')
    maximum = max(values)
    weights = [math.exp((x - maximum) / temperature) for x in values]
    ranked = sorted(range(len(values)), key=lambda i: (-values[i], i))
    candidates = ranked if top_k is None else ranked[:top_k]
    if top_p is not None and top_p < 1:
        mass = math.fsum(weights[i] for i in candidates)
        cumulative = 0.0
        kept = []
        for i in candidates:
            kept.append(i)
            cumulative += weights[i] / mass
            if cumulative >= top_p:
                break
        candidates = kept
    total = math.fsum(weights[i] for i in candidates)
    selected = set(candidates)
    return [weight / total if i in selected else 0.0 for i, weight in enumerate(weights)]


def sampling_scenario():
    """Compute four policies on the same illustrative logits, with filters applied separately."""
    logits = [3.0, 2.0, 1.0, 0.0]
    return {
        'logits': logits,
        'temperature_1': token_distribution(logits),
        'temperature_2': token_distribution(logits, temperature=2),
        'top_k_2': token_distribution(logits, top_k=2),
        'top_p_09': token_distribution(logits, top_p=0.9),
    }

scenario = sampling_scenario()
for name, values in scenario.items():
    print(name, [round(value, 6) for value in values])
fig = sampling_plot(scenario, 'en')
plt.show()
plt.close(fig)

import random

z = torch.tensor(scenario['logits'], dtype=torch.float64)
expected = torch.softmax(z, dim=0)
torch.testing.assert_close(torch.tensor(scenario['temperature_1'], dtype=torch.float64), expected)
torch.testing.assert_close(torch.tensor(scenario['temperature_2'], dtype=torch.float64), torch.softmax(z / 2, dim=0))
for policy, count in [('top_k_2', 2), ('top_p_09', 3)]:
    truncated = torch.zeros_like(z)
    truncated[:count] = torch.softmax(z[:count], dim=0)
    torch.testing.assert_close(torch.tensor(scenario[policy], dtype=torch.float64), truncated)
assert token_distribution([0., 0., 0., 0.], top_p=0.5) == [0.5, 0.5, 0., 0.]
assert token_distribution([0., 0.], top_k=1) == [1., 0.]

probabilities = scenario['top_p_09']
draws = random.Random(7).choices(range(4), weights=probabilities, k=20_000)
frequencies = [draws.count(i) / len(draws) for i in range(4)]
assert max(abs(a - b) for a, b in zip(frequencies, probabilities)) < 0.015
print({'probabilities': probabilities, 'frequencies': frequencies})

def greedy_trace(prepared_logits, eos_id, max_new_tokens):
    if isinstance(max_new_tokens, bool) or not isinstance(max_new_tokens, Integral) or max_new_tokens <= 0:
        raise ValueError('max_new_tokens must be a positive integer')
    if isinstance(eos_id, bool) or not isinstance(eos_id, Integral) or eos_id < 0:
        raise ValueError('eos_id must be a nonnegative integer')
    generated = []
    for scores in prepared_logits:
        probabilities = token_distribution(scores)
        if eos_id >= len(probabilities):
            raise ValueError('eos_id is outside the score vector')
        token = max(range(len(probabilities)), key=probabilities.__getitem__)
        generated.append(token)
        if token == eos_id:
            return generated, 'eos'
        if len(generated) == max_new_tokens:
            return generated, 'length_limit'
    raise ValueError('prepared trace ended before a stopping condition')

eos_result = greedy_trace([[3., 2., 1., 0.], [0., 0., 0., 3.]], eos_id=3, max_new_tokens=3)
limit_result = greedy_trace([[3., 2., 1., 0.]] * 3, eos_id=3, max_new_tokens=3)
assert eos_result == ([0, 3], 'eos')
assert limit_result == ([0, 0, 0], 'length_limit')
print(eos_result, limit_result)