Module 6. Normalization, activations, and the anatomy of a Transformer block

# Install only missing dependencies; ordinary Colab already has them.
import importlib.util
import subprocess
import sys

packages = {"numpy": "numpy", "matplotlib": "matplotlib"}
missing = [pip_name for module, pip_name in packages.items() if importlib.util.find_spec(module) is None]
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])
print('Environment ready.')

import hashlib
import json
import math
from dataclasses import dataclass
from typing import Literal

import matplotlib.pyplot as plt
import numpy as np

_REFERENCE_JSON = '{\n  "schema_version": 3,\n  "normalization": {\n    "input": [\n      2.0,\n      -1.0,\n      3.0,\n      0.0\n    ],\n    "eps": 1e-06,\n    "mean": 1.0,\n    "variance": 2.5,\n    "rms_value": 1.8708289606481936,\n    "layer_norm": [\n      0.6324554055426074,\n      -1.2649108110852147,\n      1.2649108110852147,\n      -0.6324554055426074\n    ],\n    "rms_norm": [\n      1.0690448149290206,\n      -0.5345224074645103,\n      1.6035672223935309,\n      0.0\n    ]\n  },\n  "gpt_oss_swiglu": {\n    "alpha": 1.702,\n    "limit": 7.0,\n    "diagonal_inputs": [\n      -9.0,\n      -2.0,\n      0.0,\n      2.0,\n      7.0,\n      9.0\n    ],\n    "diagonal_outputs": [\n      1.2019067919351053e-05,\n      0.06434137685579186,\n      0.0,\n      5.806975869432624,\n      55.999625026426536,\n      55.999625026426536\n    ],\n    "gelu_bridge_max_abs_error": 0.020334865464460172\n  },\n  "llama_3_1_8b": {\n    "shape": {\n      "vocab_size": 128256,\n      "d_model": 4096,\n      "n_layers": 32,\n      "d_ff": 14336,\n      "n_heads": 32,\n      "n_kv_heads": 8,\n      "d_head": 128,\n      "tied_embeddings": false\n    },\n    "parameters": {\n      "attention": 41943040,\n      "mlp": 176160768,\n      "norms": 8192,\n      "matrix_block": 218103808,\n      "full_block": 218112000,\n      "block": 218112000,\n      "embedding_tables": 1050673152,\n      "final_norm": 4096,\n      "total": 8030261248\n    },\n    "matrix_block_over_d2": 13.0,\n    "full_block_over_d2": 13.00048828125\n  },\n  "depth_experiment": {\n    "seed": 5,\n    "d_model": 256,\n    "n_layers": 36,\n    "shared_input_and_weights": true,\n    "curves": {\n      "none": [\n        1.3796141627290481,\n        1.9746924376971382,\n        2.891823928468155,\n        4.235085925879102,\n        6.197191604353382,\n        8.666202046735455,\n        11.878291733194873,\n        17.314658013366586,\n        26.14813149938547,\n        36.73399756959625,\n        51.91650763364299,\n        76.54610965962036,\n        113.41290062150921,\n        167.05196004174942,\n        235.39375521289483,\n        342.8080258207679,\n        488.44488850344885,\n        677.7200397032677,\n        896.5705008241248,\n        1330.8336285262492,\n        1889.6211261598758,\n        2767.0608667178412,\n        4195.964784744771,\n        6175.251997341848,\n        8985.95044580444,\n        11722.52879779937,\n        17161.684877954263,\n        23644.35645515886,\n        33222.01600011237,\n        46046.04577591799,\n        68361.7014074496,\n        101968.45409586924,\n        142098.5000356235,\n        207522.285062087,\n        277708.4610711304,\n        412442.09423078835\n      ],\n      "pre": [\n        1.3986422872963444,\n        1.720878100647755,\n        2.0868881952733527,\n        2.394035604115548,\n        2.6433551824491723,\n        2.8970030108290343,\n        3.044486295369714,\n        3.2361518409625405,\n        3.3995105276293374,\n        3.5298256821584983,\n        3.645937296074648,\n        3.77494661147846,\n        3.960785319631705,\n        4.094045694939046,\n        4.2008642555989075,\n        4.32213947793166,\n        4.465180513285604,\n        4.610514130181021,\n        4.777034454694044,\n        4.978826586251819,\n        5.058032163045441,\n        5.266461753684323,\n        5.343237738973209,\n        5.548367921388823,\n        5.65607830933095,\n        5.751614710699051,\n        5.761174392983736,\n        5.817986082500051,\n        5.857176875189526,\n        5.8829909732405445,\n        5.897536434915857,\n        6.05298767865147,\n        6.101077359830962,\n        6.21914525399463,\n        6.233400226034084,\n        6.423482855953835\n      ],\n      "post": [\n        1.0,\n        0.9999999999999999,\n        0.9999999999999998,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        0.9999999999999999,\n        1.0,\n        1.0,\n        0.9999999999999999,\n        1.0,\n        0.9999999999999999,\n        1.0,\n        0.9999999999999999,\n        0.9999999999999999,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        0.9999999999999999,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        0.9999999999999999,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        1.0\n      ]\n    }\n  },\n  "gpt_oss_block": {\n    "hidden_size": 2880,\n    "intermediate_size": 2880,\n    "num_experts": 128,\n    "experts_per_token": 4,\n    "swiglu_alpha": 1.702,\n    "swiglu_limit": 7.0,\n    "activation_defined_in_model_code": true\n  },\n  "dyt_surrogate": {\n    "seed": 5,\n    "sample_size": 200000,\n    "alpha_grid_start": 0.3,\n    "alpha_grid_stop": 2.0,\n    "alpha_grid_steps": 171,\n    "best_alpha": 1.4300000000000002,\n    "best_mse": 0.16170374236321125\n  }\n}\n'
_REFERENCE_SHA256 = 'd6f6ce5fdca26edabcb28b5058d85f37b524304ca1d919e42919f16142abb881'
assert hashlib.sha256(_REFERENCE_JSON.encode("utf-8")).hexdigest() == _REFERENCE_SHA256
REF = json.loads(_REFERENCE_JSON)
assert REF["schema_version"] == 3

DEMO_LANGUAGE = 'en'
plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 12, "axes.titlesize": 16, "axes.titlelocation": "left", "axes.spines.top": False, "axes.spines.right": False, "axes.prop_cycle": plt.cycler(color=["#5E81AC", "#A3BE8C", "#B48EAD"]), "grid.color": "#D8DEE9"})
np.set_printoptions(precision=5, suppress=True)
print('Module-6 reference payload unpacked and verified', _REFERENCE_SHA256[:12])

def _floating_array(x: np.ndarray) -> np.ndarray:
    """Validate and return a floating NumPy array without changing its shape."""

    array = np.asarray(x)
    if not np.issubdtype(array.dtype, np.floating):
        raise TypeError("normalization expects a floating-point array")
    if array.ndim == 0 or array.shape[-1] == 0:
        raise ValueError("the normalized last axis must be non-empty")
    if not np.isfinite(array).all():
        raise ValueError("inputs must be finite")
    return array


def _statistics_dtype(dtype: np.dtype) -> np.dtype:
    """Use at least float32 for reductions while preserving float64 inputs."""

    return np.dtype(np.float32) if np.dtype(dtype).itemsize < 4 else np.dtype(dtype)


def root_mean_square(x: np.ndarray, axis: int | None = -1, keepdims: bool = False) -> np.ndarray:
    """Compute RMS in a reduction dtype of at least float32."""

    array = _floating_array(x)
    work = array.astype(_statistics_dtype(array.dtype), copy=False)
    return np.sqrt(np.mean(work * work, axis=axis, keepdims=keepdims))


def layer_norm(
    x: np.ndarray,
    gamma: np.ndarray | float = 1.0,
    beta: np.ndarray | float = 0.0,
    eps: float = 1e-6,
) -> np.ndarray:
    """Layer-normalize `[..., d_model]` over the last axis and preserve dtype."""

    if eps <= 0 or not math.isfinite(eps):
        raise ValueError("eps must be finite and positive")
    array = _floating_array(x)
    work_dtype = _statistics_dtype(array.dtype)
    work = array.astype(work_dtype, copy=False)
    mean = work.mean(axis=-1, keepdims=True)
    variance = np.mean((work - mean) ** 2, axis=-1, keepdims=True)
    result = (work - mean) / np.sqrt(variance + eps)
    for value in (gamma, beta):
        if np.shape(value) not in ((), (array.shape[-1],)) or not np.isfinite(value).all():
            raise ValueError("affine parameters must be finite scalars or [D] vectors")
    result = result * np.asarray(gamma, dtype=work_dtype) + np.asarray(beta, dtype=work_dtype)
    return result.astype(array.dtype, copy=False)


def rms_norm(
    x: np.ndarray,
    gamma: np.ndarray | float = 1.0,
    eps: float = 1e-6,
) -> np.ndarray:
    """RMS-normalize `[..., d_model]` over the last axis and preserve dtype."""

    if eps <= 0 or not math.isfinite(eps):
        raise ValueError("eps must be finite and positive")
    array = _floating_array(x)
    work_dtype = _statistics_dtype(array.dtype)
    work = array.astype(work_dtype, copy=False)
    mean_square = np.mean(work * work, axis=-1, keepdims=True)
    result = work / np.sqrt(mean_square + eps)
    if np.shape(gamma) not in ((), (array.shape[-1],)) or not np.isfinite(gamma).all():
        raise ValueError("gamma must be a finite scalar or [D] vector")
    result = result * np.asarray(gamma, dtype=work_dtype)
    return result.astype(array.dtype, copy=False)


norm_ref = REF["normalization"]
x = np.asarray(norm_ref["input"], dtype=np.float64)
ln = layer_norm(x, eps=norm_ref["eps"])
rn = rms_norm(x, eps=norm_ref["eps"])
np.testing.assert_allclose(ln, norm_ref["layer_norm"], rtol=0, atol=1e-12)
np.testing.assert_allclose(rn, norm_ref["rms_norm"], rtol=0, atol=1e-12)
assert abs(float(ln.sum())) < 1e-12 and rn[-1] == 0.0

# eps limits amplification when mean-square is tiny: here ms = (1e-3)^2.
scales = [1.0 / math.sqrt(1e-6 + eps) for eps in (1e-6, 1e-5)]
assert abs(scales[0] / scales[1] - 2.3452078799) < 1e-9

# Fixed eps breaks exact scale invariance: exercise 6.
tiny = np.array([1e-3], dtype=np.float64)
scaled_outputs = [float(rms_norm(c * tiny, eps=1e-6)[0]) for c in (1, 2)]
np.testing.assert_allclose(scaled_outputs, [1 / math.sqrt(2), 2 / math.sqrt(5)])
print('Tiny input: RMSNorm of x and 2x:', np.round(scaled_outputs, 4))

# Squaring 1e-4 in float16 underflows; promoting before the square preserves it.
x_small = np.full(4096, 1e-4, dtype=np.float16)
naive_ms = float(np.mean(x_small * x_small, dtype=np.float16))
safe_ms = float(np.mean(x_small.astype(np.float32) ** 2, dtype=np.float32))
assert naive_ms == 0.0 and safe_ms > 0.0
safe_out = rms_norm(x_small)
assert safe_out.dtype == np.float16 and np.isfinite(safe_out).all()

print("LN :", np.round(ln, 4), "sum=", round(float(ln.sum()), 8))
print("RMS:", np.round(rn, 4), "sqrt(mean(x²)+eps)=", round(norm_ref["rms_value"], 4))
print("eps gain:", [round(v, 1) for v in scales], "· fp16 ms:", naive_ms, "· fp32 ms:", f"{safe_ms:.2e}")
print('Norms, the eps ceiling, and float16→float32 statistics verified')

def stable_sigmoid(x: np.ndarray | float) -> np.ndarray:
    """Compute sigmoid without `exp` overflow on large negative inputs."""

    z = np.asarray(x, dtype=np.float64)
    return np.exp(-np.logaddexp(0.0, -z))


def gelu_exact(x: np.ndarray) -> np.ndarray:
    """Evaluate exact GELU = x * Phi(x) without a SciPy dependency."""

    z = np.asarray(x, dtype=np.float64)
    erf_values = np.fromiter(
        (math.erf(float(v) / math.sqrt(2.0)) for v in z.ravel()),
        dtype=np.float64,
        count=z.size,
    ).reshape(z.shape)
    return 0.5 * z * (1.0 + erf_values)


def silu(x: np.ndarray | float) -> np.ndarray:
    """Return SiLU/Swish: x * sigmoid(x)."""

    z = np.asarray(x, dtype=np.float64)
    return z * stable_sigmoid(z)


def gpt_oss_swiglu(
    gate: np.ndarray | float,
    up: np.ndarray | float,
    *,
    alpha: float = 1.702,
    limit: float = 7.0,
) -> np.ndarray:
    """Apply the public gpt-oss gate clamp, up clamp, and `(up + 1)` shift."""

    if limit <= 0 or not math.isfinite(limit):
        raise ValueError("limit must be finite and positive")
    if not math.isfinite(alpha):
        raise ValueError("alpha must be finite")
    g = np.minimum(np.asarray(gate, dtype=np.float64), limit)
    u = np.clip(np.asarray(up, dtype=np.float64), -limit, limit)
    return g * stable_sigmoid(alpha * g) * (u + 1.0)


act_ref = REF["gpt_oss_swiglu"]
z = np.asarray(act_ref["diagonal_inputs"], dtype=np.float64)
got = gpt_oss_swiglu(z, z, alpha=act_ref["alpha"], limit=act_ref["limit"])
np.testing.assert_allclose(got, act_ref["diagonal_outputs"], rtol=0, atol=1e-12)
assert float(got[-1]) == float(got[-2])  # z=7 and z=9 hit the same clamp plateau.
assert np.isfinite(stable_sigmoid(np.array([-1000.0, 1000.0]))).all()

grid = np.linspace(-6.0, 6.0, 4001)
bridge = grid * stable_sigmoid(1.702 * grid)
max_error = float(np.max(np.abs(gelu_exact(grid) - bridge)))
np.testing.assert_allclose(max_error, act_ref["gelu_bridge_max_abs_error"], rtol=0, atol=1e-12)

print("z:", z)
print("gpt-oss output:", np.round(got, 5))
print("max |GELU - z*sigmoid(1.702z)| =", f"{max_error:.5f}")
print('Clamp table, saturation and GELU approximation checked')

def attention_logits(q: np.ndarray, k: np.ndarray) -> np.ndarray:
    """Return scaled dot-product logits for `[T_q, d_head]` and `[T_k, d_head]`."""

    q = _floating_array(q)
    k = _floating_array(k)
    if q.ndim != 2 or k.ndim != 2 or q.shape[-1] != k.shape[-1]:
        raise ValueError("q and k must be rank-2 with the same d_head")
    dtype = np.result_type(q.dtype, k.dtype, np.float32)
    return q.astype(dtype) @ k.astype(dtype).T / math.sqrt(q.shape[-1])


def qk_rms_normalize(x: np.ndarray, eps: float = 1e-6) -> np.ndarray:
    """Normalize each query/key vector across its head dimension."""

    return rms_norm(np.asarray(x, dtype=np.float64), eps=eps)


qk_rng = np.random.default_rng(11)
q0 = qk_rng.standard_normal((5, 16))
k0 = qk_rng.standard_normal((7, 16))
scales = np.array([1.0, 10.0, 100.0])
raw = [attention_logits(scale * q0, scale * k0) for scale in scales]
normalized = [attention_logits(qk_rms_normalize(scale * q0), qk_rms_normalize(scale * k0)) for scale in scales]

# Scaling both factors by c multiplies raw logits by c^2.
np.testing.assert_allclose(raw[1], raw[0] * 100.0, rtol=1e-12, atol=1e-12)
np.testing.assert_allclose(raw[2], raw[0] * 10_000.0, rtol=1e-12, atol=1e-10)
for logits in normalized[1:]:
    np.testing.assert_allclose(logits, normalized[0], rtol=2e-6, atol=2e-6)

print("max |raw logits|:", [round(float(np.max(np.abs(v))), 2) for v in raw])
print("max |QK-normalized logits|:", [round(float(np.max(np.abs(v))), 5) for v in normalized])
print('Approximate invariance checked at ordinary input scale')

# With eps>0, scale invariance is approximate and fails near zero.
small = np.array([[1e-5, -2e-5]])
assert np.linalg.norm(rms_norm(100*small)-rms_norm(small)) > 0.5
zero = np.zeros((2,4), dtype=np.float64)
np.testing.assert_array_equal(rms_norm(zero), zero)
constant = np.full((2,4), 3.0)
np.testing.assert_array_equal(layer_norm(constant), zero)
assert np.all(rms_norm(constant) > 0.99)

# Two heads, two features each. Changing head 1 changes head 0 only with a joint denominator.
q = np.array([[1.,2.],[3.,4.]])
changed = q.copy(); changed[1] *= 10
per_head = rms_norm(q); per_head_changed = rms_norm(changed)
joint = rms_norm(q.reshape(1,-1)).reshape(2,2)
joint_changed = rms_norm(changed.reshape(1,-1)).reshape(2,2)
np.testing.assert_array_equal(per_head[0], per_head_changed[0])
assert not np.allclose(joint[0], joint_changed[0])
print('eps and normalization-axis counterexamples: PASS')

NormPlacement = Literal["pre", "peri"]


@dataclass(frozen=True)
class BlockConfig:
    """Shape contract for the educational dense GQA + SwiGLU block."""

    d_model: int
    d_ff: int
    n_heads: int
    n_kv_heads: int
    d_head: int
    eps: float = 1e-6

    def __post_init__(self) -> None:
        """Validate canonical GQA dimensions and the numerical epsilon."""

        dims=(self.d_model, self.d_ff, self.n_heads, self.n_kv_heads, self.d_head)
        if any(isinstance(v, bool) or not isinstance(v, (int, np.integer)) or v <= 0 for v in dims):
            raise ValueError("all dimensions must be positive integers")
        if self.n_heads * self.d_head != self.d_model:
            raise ValueError("canonical block requires n_heads * d_head == d_model")
        if self.n_heads % self.n_kv_heads:
            raise ValueError("n_heads must be divisible by n_kv_heads")
        if self.eps <= 0 or not math.isfinite(self.eps):
            raise ValueError("eps must be finite and positive")


def causal_allowed(n_tokens: int) -> np.ndarray:
    """Return `[T, T]` boolean mask with `True` for visible keys."""

    if isinstance(n_tokens, bool) or not isinstance(n_tokens, (int, np.integer)) or n_tokens <= 0:
        raise ValueError("n_tokens must be a positive integer")
    return np.tri(n_tokens, n_tokens, dtype=bool)


def masked_softmax(logits: np.ndarray, allowed: np.ndarray) -> np.ndarray:
    """Row-wise softmax with `True = allowed`; reject fully masked rows."""

    scores = np.asarray(logits, dtype=np.float64)
    visible = np.broadcast_to(np.asarray(allowed, dtype=bool), scores.shape)
    if np.any(~np.any(visible, axis=-1)):
        raise ValueError("every query row must expose at least one key")
    if not np.isfinite(scores[visible]).all():
        raise FloatingPointError("allowed attention logits must be finite")
    masked = np.where(visible, scores, -np.inf)
    row_max = np.max(masked, axis=-1, keepdims=True)
    numerators = np.where(visible, np.exp(masked - row_max), 0.0)
    return numerators / np.sum(numerators, axis=-1, keepdims=True)


def sdpa(q: np.ndarray, k: np.ndarray, v: np.ndarray, allowed: np.ndarray) -> np.ndarray:
    """Reference attention for q/k/v shaped `[H, T, d_head]`."""

    q = np.asarray(q, dtype=np.float64)
    k = np.asarray(k, dtype=np.float64)
    v = np.asarray(v, dtype=np.float64)
    if q.ndim != 3 or k.shape != q.shape or v.shape != q.shape:
        raise ValueError("q, k, and v must share shape [H, T, d_head]")
    logits = q @ np.swapaxes(k, -1, -2) / math.sqrt(q.shape[-1])
    return masked_softmax(logits, allowed) @ v


def grouped_query_attention(
    q_flat: np.ndarray,
    k_flat: np.ndarray,
    v_flat: np.ndarray,
    config: BlockConfig,
    allowed: np.ndarray,
) -> np.ndarray:
    """Apply GQA to flat projections and return `[T, d_model]`."""

    q_flat = np.asarray(q_flat)
    k_flat = np.asarray(k_flat)
    v_flat = np.asarray(v_flat)
    if q_flat.ndim != 2 or k_flat.ndim != 2 or v_flat.ndim != 2:
        raise ValueError("q_flat, k_flat, and v_flat must be rank-2")
    n_tokens = q_flat.shape[0]
    allowed = np.asarray(allowed, dtype=bool)
    if allowed.shape != (n_tokens, n_tokens):
        raise ValueError("allowed must have shape [T, T]")
    expected = (
        (n_tokens, config.d_model),
        (n_tokens, config.n_kv_heads * config.d_head),
        (n_tokens, config.n_kv_heads * config.d_head),
    )
    if (q_flat.shape, k_flat.shape, v_flat.shape) != expected:
        raise ValueError(f"projection shapes must be {expected}")
    q = q_flat.reshape(n_tokens, config.n_heads, config.d_head).transpose(1, 0, 2)
    k = k_flat.reshape(n_tokens, config.n_kv_heads, config.d_head).transpose(1, 0, 2)
    v = v_flat.reshape(n_tokens, config.n_kv_heads, config.d_head).transpose(1, 0, 2)
    repeats = config.n_heads // config.n_kv_heads
    k = np.repeat(k, repeats, axis=0)
    v = np.repeat(v, repeats, axis=0)
    return sdpa(q, k, v, allowed).transpose(1, 0, 2).reshape(n_tokens, config.d_model)


class PreNormBlock:
    """Small NumPy block that exposes both residual updates for inspection."""

    def __init__(self, config: BlockConfig, rng: np.random.Generator):
        """Sample one deterministic bias-free block from ``rng``."""

        self.config = config

        def weight(fan_in: int, fan_out: int) -> np.ndarray:
            """Return a fan-in-scaled matrix shaped ``[fan_in, fan_out]``."""

            return rng.standard_normal((fan_in, fan_out)) / math.sqrt(fan_in)

        d, i = config.d_model, config.d_ff
        q_width = config.n_heads * config.d_head
        kv_width = config.n_kv_heads * config.d_head
        self.gamma_attn = np.ones(d)
        self.gamma_ffn = np.ones(d)
        self.wq = weight(d, q_width)
        self.wk = weight(d, kv_width)
        self.wv = weight(d, kv_width)
        self.wo = weight(q_width, d)
        self.w_gate = weight(d, i)
        self.w_up = weight(d, i)
        self.w_down = weight(i, d)

    def attention_update(self, x: np.ndarray, allowed: np.ndarray) -> np.ndarray:
        """Return the attention branch update before residual addition."""

        z = rms_norm(x, self.gamma_attn, self.config.eps)
        mixed = grouped_query_attention(
            z @ self.wq,
            z @ self.wk,
            z @ self.wv,
            self.config,
            allowed,
        )
        return mixed @ self.wo

    def ffn_update(self, x: np.ndarray) -> np.ndarray:
        """Return the standard SwiGLU FFN update before residual addition."""

        z = rms_norm(x, self.gamma_ffn, self.config.eps)
        gate = z @ self.w_gate
        up = z @ self.w_up
        return (silu(gate) * up) @ self.w_down

    def forward(
        self,
        x: np.ndarray,
        allowed: np.ndarray,
        *,
        placement: NormPlacement = "pre",
        enable_attention: bool = True,
        enable_ffn: bool = True,
    ) -> tuple[np.ndarray, dict[str, np.ndarray]]:
        """Run pre-norm or Peri-LN and return output plus both updates."""

        if placement not in ("pre", "peri"):
            raise ValueError("placement must be 'pre' or 'peri'")
        h = np.asarray(x, dtype=np.float64)
        if h.ndim != 2 or h.shape[-1] != self.config.d_model:
            raise ValueError("x must have shape [T, d_model]")
        if not np.isfinite(h).all():
            raise FloatingPointError("x must contain only finite values")

        attn_update = np.zeros_like(h)
        if enable_attention:
            attn_update = self.attention_update(h, allowed)
            if placement == "peri":
                attn_update = rms_norm(attn_update, eps=self.config.eps)
        after_attention = h + attn_update

        ffn_update = np.zeros_like(h)
        if enable_ffn:
            ffn_update = self.ffn_update(after_attention)
            if placement == "peri":
                ffn_update = rms_norm(ffn_update, eps=self.config.eps)
        output = after_attention + ffn_update
        return output, {
            "attention_update": attn_update,
            "after_attention": after_attention,
            "ffn_update": ffn_update,
        }


block_rng = np.random.default_rng(5)
config = BlockConfig(d_model=128, d_ff=448, n_heads=8, n_kv_heads=2, d_head=16)
block = PreNormBlock(config, block_rng)
x_block = block_rng.standard_normal((12, config.d_model))
allowed = causal_allowed(len(x_block))

output, details = block.forward(x_block, allowed)
identity, identity_details = block.forward(
    x_block, allowed, enable_attention=False, enable_ffn=False
)
no_ffn, no_ffn_details = block.forward(x_block, allowed, enable_ffn=False)

assert output.shape == x_block.shape and np.isfinite(output).all()
np.testing.assert_array_equal(identity, x_block)
np.testing.assert_array_equal(identity_details["attention_update"], np.zeros_like(x_block))
np.testing.assert_allclose(no_ffn, x_block + no_ffn_details["attention_update"], rtol=0, atol=1e-12)
np.testing.assert_allclose(output - no_ffn, details["ffn_update"], rtol=0, atol=1e-12)

changed = x_block.copy()
changed[6:] += 25 * block_rng.standard_normal(changed[6:].shape)
changed_out, _ = block.forward(changed, allowed)
np.testing.assert_allclose(changed_out[:6], output[:6], rtol=0, atol=1e-12)
assert not np.allclose(changed_out[6:], output[6:])

# Broken masks and non-finite residual states fail loudly instead of producing NaN.
for broken_call, expected_error in (
    (lambda: masked_softmax(np.zeros((2, 2)), np.zeros((2, 2), dtype=bool)), ValueError),
    (lambda: block.forward(np.full_like(x_block, np.inf), allowed), FloatingPointError),
    (lambda: BlockConfig(128, 448, 8, 2, 16, eps=0.0), ValueError),
):
    try:
        broken_call()
    except expected_error:
        pass
    else:
        raise AssertionError(f"expected {expected_error.__name__}")

print("output:", output.shape, "· attention update rms:", round(float(root_mean_square(details["attention_update"], axis=None)), 4))
print("ffn update rms:", round(float(root_mean_square(details["ffn_update"], axis=None)), 4))
print('Shapes, residual updates and causality checked')

@dataclass(frozen=True)
class DecoderShape:
    """Fields needed to count a bias-free dense decoder-only Transformer."""

    vocab_size: int
    d_model: int
    n_layers: int
    d_ff: int
    n_heads: int
    n_kv_heads: int
    d_head: int
    tied_embeddings: bool = False

    def __post_init__(self) -> None:
        """Validate the dimensions assumed by the exact parameter counter."""

        dims=(
            self.vocab_size,
            self.d_model,
            self.n_layers,
            self.d_ff,
            self.n_heads,
            self.n_kv_heads,
            self.d_head,
        )
        if any(isinstance(v, bool) or not isinstance(v, (int, np.integer)) or v <= 0 for v in dims):
            raise ValueError("all shape fields must be positive integers")
        if not isinstance(self.tied_embeddings, bool):
            raise ValueError("tied_embeddings must be boolean")
        if self.n_heads * self.d_head != self.d_model:
            raise ValueError("n_heads * d_head must equal d_model for this counter")
        if self.n_heads % self.n_kv_heads:
            raise ValueError("n_heads must be divisible by n_kv_heads")


@dataclass(frozen=True)
class ParameterBreakdown:
    """Exact count split for the bias-free decoder assumed below."""

    attention: int
    mlp: int
    norms: int
    matrix_block: int
    full_block: int
    block: int  # backwards-compatible alias of full_block in the frozen contract
    embedding_tables: int
    final_norm: int
    total: int


def count_decoder_parameters(shape: DecoderShape) -> ParameterBreakdown:
    """Count GQA + SwiGLU weights, norm scales, embeddings, and final norm."""

    q_width = shape.n_heads * shape.d_head
    kv_width = shape.n_kv_heads * shape.d_head
    attention = shape.d_model * q_width + q_width * shape.d_model
    attention += 2 * shape.d_model * kv_width
    mlp = 3 * shape.d_model * shape.d_ff
    norms = 2 * shape.d_model
    matrix_block = attention + mlp
    full_block = matrix_block + norms
    tables = (1 if shape.tied_embeddings else 2) * shape.vocab_size * shape.d_model
    final_norm = shape.d_model
    total = tables + shape.n_layers * full_block + final_norm
    return ParameterBreakdown(
        attention, mlp, norms, matrix_block, full_block, full_block, tables, final_norm, total
    )


small_shape = DecoderShape(128, config.d_model, 1, config.d_ff, config.n_heads, config.n_kv_heads, config.d_head)
actual_block = sum(v.size for v in vars(block).values() if isinstance(v, np.ndarray))
assert actual_block == count_decoder_parameters(small_shape).full_block

# Count each distinct stored array once; aliasing the head implements weight tying.
block_arrays = [v for v in vars(block).values() if isinstance(v, np.ndarray)]
embeddings = np.zeros((small_shape.vocab_size, config.d_model))
final_gain = np.ones(config.d_model)
for tied in (False, True):
    head = embeddings if tied else np.zeros_like(embeddings)
    arrays = block_arrays + [embeddings, final_gain, head]
    unique_arrays = {id(array): array for array in arrays}
    actual_total = sum(array.size for array in unique_arrays.values())
    shape_fields = dict(vars(small_shape), tied_embeddings=tied)
    assert actual_total == count_decoder_parameters(DecoderShape(**shape_fields)).total

shape = DecoderShape(**REF["llama_3_1_8b"]["shape"])
breakdown = count_decoder_parameters(shape)
expected = REF["llama_3_1_8b"]["parameters"]
assert breakdown.__dict__ == expected
assert breakdown.total == 8_030_261_248
assert breakdown.matrix_block == 13 * shape.d_model**2
assert breakdown.full_block == breakdown.matrix_block + 2 * shape.d_model
np.testing.assert_allclose(
    breakdown.full_block / shape.d_model**2,
    REF["llama_3_1_8b"]["full_block_over_d2"],
    rtol=0, atol=1e-12,
)

try:
    DecoderShape(128, 64, 2, 128, 3, 1, 16)
except ValueError:
    pass
else:
    raise AssertionError("incompatible n_heads*d_head must be rejected")

print("matrix block:", f"{breakdown.matrix_block/1e6:.3f}M", "· full block:", f"{breakdown.full_block/1e6:.3f}M")
print("attention:", f"{breakdown.attention/breakdown.full_block:.1%}", "· FFN:", f"{breakdown.mlp/breakdown.full_block:.1%}")
print("model:", f"{breakdown.total:,}", "parameters")
print('Exact count for the stated 8B geometry checked')

DepthMode = Literal["none", "pre", "post"]


def scalar_rms(x: np.ndarray) -> float:
    """Return float64 RMS for a scalar diagnostic."""

    z = np.asarray(x, dtype=np.float64)
    return float(np.sqrt(np.mean(z * z)))


def unit_rms(x: np.ndarray, floor: float = 1e-12) -> np.ndarray:
    """Normalize to unit RMS; `floor` only guards a degenerate zero vector."""

    if floor <= 0 or not math.isfinite(floor):
        raise ValueError("floor must be finite and positive")
    return np.asarray(x, dtype=np.float64) / max(scalar_rms(x), floor)


def make_depth_problem(seed: int, d_model: int, n_layers: int) -> tuple[np.ndarray, list[np.ndarray]]:
    """Generate one input and one weight stack shared by all modes."""

    if d_model <= 0 or n_layers <= 0:
        raise ValueError("d_model and n_layers must be positive")
    rng = np.random.default_rng(seed)
    x0 = rng.standard_normal(d_model)
    weights = [
        rng.standard_normal((d_model, d_model)) / math.sqrt(d_model)
        for _ in range(n_layers)
    ]
    return x0, weights


def simulate_depth(mode: DepthMode, x0: np.ndarray, weights: list[np.ndarray]) -> list[float]:
    """Run none/pre/post placement on the same linear residual stack."""

    if mode not in ("none", "pre", "post"):
        raise ValueError("unknown mode")
    h = np.asarray(x0, dtype=np.float64).copy()
    if h.ndim != 1 or not weights:
        raise ValueError("x0 must be a vector and weights must be non-empty")
    curve = []
    for weight in weights:
        weight = np.asarray(weight, dtype=np.float64)
        if weight.shape != (h.size, h.size):
            raise ValueError("every weight must have shape [d_model, d_model]")
        if mode == "none":
            h = h + weight @ h
        elif mode == "pre":
            h = h + weight @ unit_rms(h)
        else:
            h = unit_rms(h + weight @ h)
        curve.append(scalar_rms(h))
    return curve


depth_ref = REF["depth_experiment"]
x0, weights = make_depth_problem(
    depth_ref["seed"], depth_ref["d_model"], depth_ref["n_layers"]
)
curves = {mode: simulate_depth(mode, x0, weights) for mode in ("none", "pre", "post")}
assert depth_ref["shared_input_and_weights"] is True
for mode, curve in curves.items():
    np.testing.assert_allclose(
        curve,
        depth_ref["curves"][mode],
        rtol=1e-8,
        atol=1e-10,
    )
assert curves["none"][-1] > 4e5
assert 6.0 < curves["pre"][-1] < 7.0
np.testing.assert_allclose(curves["post"], 1.0, rtol=0, atol=1e-12)

fig, ax = plt.subplots(figsize=(11, 5.4), dpi=150)
for mode, curve in curves.items():
    ax.plot(range(1, len(curve) + 1), curve, label=mode)
ax.set_yscale("log")
ax.set_xlabel("Слой" if DEMO_LANGUAGE == "ru" else "Layer")
ax.set_ylabel("RMS состояния · лог. шкала" if DEMO_LANGUAGE == "ru" else "Residual RMS · log scale")
ax.set_title("Один вход и веса: масштаб по глубине" if DEMO_LANGUAGE == "ru" else "Shared input and weights: scale with depth")
ax.grid(alpha=0.3)
ax.legend()
fig.tight_layout()
plt.show()
plt.close(fig)

print("layer 36:", {mode: round(curve[-1], 3) for mode, curve in curves.items()})
print('Shared-input depth experiment reproduced')

def make_block_stack(
    seed: int,
    n_layers: int,
    config: BlockConfig,
) -> list[PreNormBlock]:
    """Create one deterministic stack reused by pre and Peri-LN traces."""

    rng = np.random.default_rng(seed)
    return [PreNormBlock(config, rng) for _ in range(n_layers)]


def trace_block_stack(
    blocks: list[PreNormBlock],
    x0: np.ndarray,
    allowed: np.ndarray,
    placement: NormPlacement,
) -> dict[str, list[float]]:
    """Record residual, attention-update, and FFN-update RMS by layer."""

    x = np.asarray(x0, dtype=np.float64).copy()
    trace = {"residual": [], "attention_update": [], "ffn_update": []}
    for block in blocks:
        x, details = block.forward(x, allowed, placement=placement)
        trace["residual"].append(scalar_rms(x))
        trace["attention_update"].append(scalar_rms(details["attention_update"]))
        trace["ffn_update"].append(scalar_rms(details["ffn_update"]))
    return trace


stack_config = BlockConfig(d_model=128, d_ff=448, n_heads=8, n_kv_heads=2, d_head=16)
blocks = make_block_stack(seed=7, n_layers=24, config=stack_config)
input_rng = np.random.default_rng(70)
x0_stack = input_rng.standard_normal((8, stack_config.d_model))
stack_mask = causal_allowed(len(x0_stack))

pre_trace = trace_block_stack(blocks, x0_stack, stack_mask, "pre")
peri_trace = trace_block_stack(blocks, x0_stack, stack_mask, "peri")

for branch in ("attention_update", "ffn_update"):
    np.testing.assert_allclose(peri_trace[branch], 1.0, rtol=0, atol=2e-5)
assert np.std(pre_trace["attention_update"]) > 0.01
assert np.std(pre_trace["ffn_update"]) > 0.01
assert peri_trace["residual"][-1] > 1.0

fig, ax = plt.subplots(figsize=(11, 5.4), dpi=150)
ax.plot(pre_trace["residual"], label="Pre-LN")
ax.plot(peri_trace["residual"], label="Peri-LN")
ax.set_xlabel("Слой" if DEMO_LANGUAGE == "ru" else "Layer")
ax.set_ylabel("RMS")
ax.set_title("Один стек блоков, два размещения нормы" if DEMO_LANGUAGE == "ru" else "One frozen stack, two norm placements")
ax.grid(alpha=0.3)
ax.legend()
fig.tight_layout()
plt.show()
plt.close(fig)

print("residual RMS at layer 24 · pre:", round(pre_trace["residual"][-1], 3), "· peri:", round(peri_trace["residual"][-1], 3))
print("Peri branch RMS ranges:",
      (min(peri_trace["attention_update"]), max(peri_trace["attention_update"])),
      (min(peri_trace["ffn_update"]), max(peri_trace["ffn_update"])))
print('Peri-LN normalized both updates while the residual stream still accumulated')

def dynamic_tanh(
    x: np.ndarray,
    alpha: float,
    gamma: np.ndarray | float = 1.0,
    beta: np.ndarray | float = 0.0,
) -> np.ndarray:
    """Element-wise DyT-style transform `gamma*tanh(alpha*x) + beta`."""

    if not math.isfinite(alpha):
        raise ValueError("alpha must be finite")
    z = np.asarray(x, dtype=np.float64)
    return np.asarray(gamma) * np.tanh(alpha * z) + np.asarray(beta)


def dyisru(x: np.ndarray, b: float, gamma: np.ndarray | float = 1.0, beta: np.ndarray | float = 0.0) -> np.ndarray:
    """DyISRU: gamma*x/sqrt(b+x**2)+beta, with b>0; sqrt(D) may be absorbed in gamma."""

    if not math.isfinite(b) or b <= 0:
        raise ValueError("b must be finite and positive")
    z = np.asarray(x, dtype=np.float64)
    return np.asarray(gamma) * z / np.hypot(math.sqrt(b), z) + np.asarray(beta)


# Same first coordinate, different neighbors: RMSNorm couples coordinates.
pair = np.array([[1.0, 1.0], [1.0, 10.0]])
rms_first = rms_norm(pair)[:, 0]
dyt_first = dynamic_tanh(pair, alpha=0.5)[:, 0]
dyisru_first = dyisru(pair, b=4.0)[:, 0]
assert not np.isclose(rms_first[0], rms_first[1])
np.testing.assert_allclose(dyt_first[0], dyt_first[1], rtol=0, atol=0)
np.testing.assert_allclose(dyisru_first[0], dyisru_first[1], rtol=0, atol=0)

surrogate = REF["dyt_surrogate"]
sample_rng = np.random.default_rng(surrogate["seed"])
sample = sample_rng.standard_normal(surrogate["sample_size"])
target = sample / scalar_rms(sample)
alpha_grid = np.linspace(
    surrogate["alpha_grid_start"],
    surrogate["alpha_grid_stop"],
    surrogate["alpha_grid_steps"],
)
mse = np.array([np.mean((np.tanh(alpha * sample) - target) ** 2) for alpha in alpha_grid])
best_alpha = float(alpha_grid[int(np.argmin(mse))])
np.testing.assert_allclose(best_alpha, surrogate["best_alpha"], rtol=0, atol=1e-12)
np.testing.assert_allclose(float(mse.min()), surrogate["best_mse"], rtol=1e-12, atol=1e-12)

grid = np.linspace(-4.5, 4.5, 901)
fig, ax = plt.subplots(figsize=(11, 5.4), dpi=150)
ax.plot(grid, grid, linestyle=":", label="y = x")
ax.plot(grid, dynamic_tanh(grid, 0.5), label="DyT α=0.5")
ax.plot(grid, dyisru(grid, 4.0), label="DyISRU b=4")
ax.set_ylim(-2.5, 2.5)
ax.set_xlabel("x")
ax.set_ylabel("Выход" if DEMO_LANGUAGE == "ru" else "Output")
ax.grid(alpha=0.3)
ax.legend()
fig.tight_layout()
plt.show()
plt.close(fig)

print("first-coordinate outputs · RMSNorm:", np.round(rms_first, 4), "· DyT:", np.round(dyt_first, 4))
print("best scalar-surrogate alpha:", best_alpha, "· MSE:", round(float(mse.min()), 6))
print('Element-wise contract and alpha≈1.43 verified')

def sinkhorn_mix(logits: np.ndarray, iterations: int = 20) -> np.ndarray:
    """Approximate a doubly stochastic matrix; callers must check both residuals."""
    log_a = np.asarray(logits, dtype=np.float64).copy()
    if log_a.ndim != 2 or log_a.shape[0] != log_a.shape[1] or not log_a.size or not np.isfinite(log_a).all():
        raise ValueError("logits must be a finite, nonempty square matrix")
    if isinstance(iterations, bool) or not isinstance(iterations, (int, np.integer)) or iterations < 1:
        raise ValueError("iterations must be a positive integer")
    # Alternate row/column normalization in log space to avoid exp overflow.
    for _ in range(iterations):
        log_a -= np.logaddexp.reduce(log_a, axis=1, keepdims=True)
        log_a -= np.logaddexp.reduce(log_a, axis=0, keepdims=True)
    return np.exp(log_a)

logits = np.array([[5.,1.,-2.],[0.,3.,-1.],[2.,0.,1.]])
records = []
for steps in (1,20,100):
    a = sinkhorn_mix(logits, steps)
    row_error = float(np.max(np.abs(a.sum(1)-1)))
    column_error = float(np.max(np.abs(a.sum(0)-1)))
    records.append((steps,row_error,column_error))
a = sinkhorn_mix(logits,100)
np.testing.assert_allclose(a.sum(0),1,atol=1e-10)
np.testing.assert_allclose(a.sum(1),1,atol=1e-10)
assert np.all(a>=0)
streams = np.array([[2.,-3.],[4.,7.],[-1.,5.]])
mixed = a @ streams
assert np.all(mixed <= streams.max(0)+1e-9)
assert np.all(mixed >= streams.min(0)-1e-9)
assert not np.allclose(mixed,streams)  # Doubly stochastic does not mean identity.
# This tests a fixed mixing matrix, not a full mHC layer or its input-dependent Jacobian.
print('Sinkhorn (iterations, row error, column error):', records)

import importlib.util
import subprocess
import sys

missing = []
if importlib.util.find_spec("numpy") is None:
    missing.append("numpy>=1.26")
if importlib.util.find_spec("torch") is None:
    missing.append("torch>=2.5")
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])

import numpy as np
import torch
import torch.nn.functional as F
import math

def numpy_rms_norm(x: np.ndarray, weight: np.ndarray | None = None, *, eps: float = 1e-6) -> np.ndarray:
    """RMSNorm over the last dimension using float64 reference reductions."""
    x = np.asarray(x)
    if x.ndim == 0 or x.shape[-1] == 0 or not np.isfinite(x).all() or not np.isfinite(eps) or eps <= 0:
        raise ValueError("x must be finite and eps positive")
    work = x.astype(np.float64, copy=False)
    y = work / np.sqrt(np.mean(work * work, axis=-1, keepdims=True) + eps)
    if weight is not None:
        w = np.asarray(weight, dtype=np.float64)
        if w.shape != (x.shape[-1],) or not np.isfinite(w).all():
            raise ValueError("weight must be finite with shape [D]")
        y = y * w
    return y.astype(x.dtype, copy=False) if np.issubdtype(x.dtype, np.floating) else y

def torch_rms_norm_explicit(
    x: torch.Tensor,
    weight: torch.Tensor | None = None,
    *,
    eps: float = 1e-6,
) -> torch.Tensor:
    """Visible tensor RMSNorm; reduction is written explicitly for teaching."""
    if not x.is_floating_point() or x.ndim == 0 or x.shape[-1] == 0:
        raise ValueError("x must be floating with a nonempty last axis")
    if not np.isfinite(eps) or eps <= 0:
        raise ValueError("eps must be finite and positive")
    work = x.to(torch.float32) if x.dtype in (torch.float16, torch.bfloat16) else x
    rms_inv = torch.rsqrt(work.square().mean(dim=-1, keepdim=True) + eps)
    y = work * rms_inv
    if weight is not None:
        if weight.shape != (x.shape[-1],):
            raise ValueError("weight must have shape [D]")
        y = y * weight.to(work.dtype)
    return y.to(x.dtype)

def torch_rms_norm(x: torch.Tensor, weight: torch.Tensor | None = None, *, eps: float = 1e-6) -> torch.Tensor:
    """PyTorch RMSNorm bridge with native autograd and device handling."""
    return F.rms_norm(x, [x.shape[-1]], weight=weight, eps=eps)

rng = np.random.default_rng(7)
x_np = rng.standard_normal((3, 5)); w_np = rng.standard_normal(5)
np_out = numpy_rms_norm(x_np, w_np, eps=1e-6)
x = torch.tensor(x_np, dtype=torch.float64, requires_grad=True)
w = torch.tensor(w_np, dtype=torch.float64, requires_grad=True)
explicit = torch_rms_norm_explicit(x, w, eps=1e-6)
api = torch_rms_norm(x, w, eps=1e-6)
np.testing.assert_allclose(np_out, explicit.detach().numpy(), rtol=1e-10, atol=1e-10)
torch.testing.assert_close(api, explicit, rtol=1e-10, atol=1e-10)
assert torch.autograd.gradcheck(lambda a, b: torch_rms_norm_explicit(a, b, eps=1e-6), (x, w))
for a, b in zip(torch.autograd.grad(explicit.square().sum(), (x, w)),
                torch.autograd.grad(api.square().sum(), (x, w))):
    torch.testing.assert_close(a, b, rtol=1e-9, atol=1e-10)
for dtype in (torch.float16, torch.bfloat16):
    large = torch.full((2, 8), 1000.0, dtype=dtype)
    torch.testing.assert_close(torch_rms_norm_explicit(large), torch.ones_like(large))
print("B07 explicit NumPy / PyTorch / F.rms_norm: PASS")

import importlib.util
import subprocess
import sys

missing = []
if importlib.util.find_spec("numpy") is None:
    missing.append("numpy>=1.26")
if importlib.util.find_spec("torch") is None:
    missing.append("torch>=2.5")
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])

import numpy as np
import torch
import torch.nn.functional as F
import math

def numpy_swiglu(
    gate: np.ndarray,
    up: np.ndarray,
    *,
    alpha: float = 1.0,
    gate_limit: float | None = None,
    up_limit: float | None = None,
    shift_up: float = 0.0,
) -> np.ndarray:
    """SwiGLU; optional limits and shift reproduce the gpt-oss-style teaching variant."""
    gate, up = np.asarray(gate), np.asarray(up)
    if gate.shape != up.shape or not np.isfinite(gate).all() or not np.isfinite(up).all():
        raise ValueError("gate and up must be finite arrays of the same shape")
    if not np.isfinite(alpha) or not np.isfinite(shift_up):
        raise ValueError("alpha and shift_up must be finite")
    if any(v is not None and (not np.isfinite(v) or v <= 0) for v in (gate_limit, up_limit)):
        raise ValueError("limits must be finite and positive")
    g = gate.astype(np.float64, copy=False)
    u = up.astype(np.float64, copy=False)
    if gate_limit is not None:
        g = np.minimum(g, gate_limit)
    if up_limit is not None:
        u = np.clip(u, -up_limit, up_limit)
    sigmoid = np.exp(-np.logaddexp(0.0, -alpha * g))
    y = g * sigmoid * (u + shift_up)
    dtype = np.result_type(gate.dtype, up.dtype)
    return y.astype(dtype, copy=False) if np.issubdtype(dtype, np.floating) else y

def torch_swiglu(
    gate: torch.Tensor,
    up: torch.Tensor,
    *,
    alpha: float = 1.0,
    gate_limit: float | None = None,
    up_limit: float | None = None,
    shift_up: float = 0.0,
) -> torch.Tensor:
    """Tensor equivalent of :func:`numpy_swiglu`, preserving autograd."""
    if gate.shape != up.shape:
        raise ValueError("gate and up must have the same shape")
    if gate.dtype != up.dtype or not gate.is_floating_point() or gate.device != up.device:
        raise ValueError("gate and up must share a floating dtype and device")
    if not np.isfinite(alpha) or not np.isfinite(shift_up):
        raise ValueError("alpha and shift_up must be finite")
    if any(v is not None and (not np.isfinite(v) or v <= 0) for v in (gate_limit, up_limit)):
        raise ValueError("limits must be finite and positive")
    dtype = torch.float32 if gate.dtype in (torch.float16, torch.bfloat16) else gate.dtype
    g, u = gate.to(dtype), up.to(dtype)
    if gate_limit is not None:
        g = g.clamp(max=gate_limit)
    if up_limit is not None:
        u = u.clamp(min=-up_limit, max=up_limit)
    return (g * torch.sigmoid(alpha * g) * (u + shift_up)).to(gate.dtype)

rng = np.random.default_rng(8)
g_np = rng.uniform(-4, 4, (3, 5)); u_np = rng.uniform(-4, 4, (3, 5))
kwargs = dict(alpha=1.702, gate_limit=7.0, up_limit=7.0, shift_up=1.0)
np_out = numpy_swiglu(g_np, u_np, **kwargs)
g = torch.tensor(g_np, dtype=torch.float64, requires_grad=True); u = torch.tensor(u_np, dtype=torch.float64, requires_grad=True)
t_out = torch_swiglu(g, u, **kwargs)
np.testing.assert_allclose(np_out, t_out.detach().numpy(), rtol=1e-10, atol=1e-10)
assert torch.autograd.gradcheck(lambda a, b: torch_swiglu(a, b, **kwargs), (g, u))
extreme = np.array([-1000., -9., 0., 9., 1000.])
with np.errstate(over='raise', invalid='raise'):
    stable = numpy_swiglu(extreme, np.ones_like(extreme), **kwargs)
np.testing.assert_allclose(stable, torch_swiglu(torch.tensor(extreme), torch.ones(5, dtype=torch.float64), **kwargs))
outside = torch.tensor([-9., 9.], dtype=torch.float64, requires_grad=True)
assert torch.autograd.gradcheck(lambda a: torch_swiglu(a, a, **kwargs), (outside,))
print("B08 forward + gradient: PASS")

from __future__ import annotations
import importlib.util, subprocess, sys
for package in ("numpy", "matplotlib"):
    if importlib.util.find_spec(package) is None:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])


import math
import random
import time
from collections import Counter, defaultdict

import matplotlib.pyplot as plt
import numpy as np

SEED = 20260801
random.seed(SEED)
np.random.seed(SEED)
print(f"reproducible smoke tier; seed={SEED}")
DEMO_LANGUAGE = 'en'

if importlib.util.find_spec("torch") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "torch"])

import torch
from torch import nn
import torch.nn.functional as F
torch.set_num_threads(1)
ru = DEMO_LANGUAGE == "ru"
plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 12,
    "axes.spines.top": False, "axes.spines.right": False,
    "axes.prop_cycle": plt.cycler(color=["#5E81AC", "#A3BE8C", "#B48EAD"])})
data_rng = torch.Generator().manual_seed(SEED)
x = torch.randn(768, 32, generator=data_rng) * torch.linspace(.5, 3., 32)
y = torch.sin(x[:, 0]) + .2*x[:, 1]
x_train, x_valid = x[:512], x[512:]
y_train, y_valid = y[:512], y[512:]

class RMSNorm(nn.Module):
    def __init__(self):
        super().__init__()
        self.gain = nn.Parameter(torch.ones(32))
    def forward(self, z):
        return self.gain * z * torch.rsqrt(z.square().mean(-1, keepdim=True)+1e-6)

class DyT(nn.Module):
    def __init__(self):
        super().__init__()
        self.alpha = nn.Parameter(torch.tensor(.5))
        self.gain = nn.Parameter(torch.ones(32))
        self.bias = nn.Parameter(torch.zeros(32))
    def forward(self, z):
        return self.gain * torch.tanh(self.alpha*z) + self.bias

class DyISRU(nn.Module):
    def __init__(self):
        super().__init__()
        # b=4 makes the initial slope 1/sqrt(b)=0.5, matching DyT.
        self.raw_b = nn.Parameter(torch.log(torch.expm1(torch.tensor(4.))))
        self.gain = nn.Parameter(torch.ones(32))
        self.bias = nn.Parameter(torch.zeros(32))
    def forward(self, z):
        b = F.softplus(self.raw_b) + 1e-8
        return self.gain * z * torch.rsqrt(b+z.square()) + self.bias

def train(norm_cls):
    torch.manual_seed(SEED)  # These normalizers consume no random draws.
    model = nn.Sequential(nn.Linear(32,32), norm_cls(), nn.SiLU(), nn.Linear(32,1))
    initial = [p.detach().clone() for i in (0,3) for p in model[i].parameters()]
    opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
    history = []
    for _ in range(80):
        model.train()
        loss = F.mse_loss(model(x_train).squeeze(-1), y_train)
        assert torch.isfinite(loss)
        opt.zero_grad(); loss.backward()
        assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters())
        opt.step()
        model.eval()
        with torch.no_grad():
            train_mse = F.mse_loss(model(x_train).squeeze(-1), y_train)
            valid = F.mse_loss(model(x_valid).squeeze(-1), y_valid)
        history.append((train_mse.item(),valid.item()))
    return np.asarray(history),initial,sum(p.numel() for p in model.parameters())

fig,axes=plt.subplots(1,2,figsize=(11,4.8),dpi=150)
initial_reference = None
for cls in (RMSNorm,DyT,DyISRU):
    history,initial,count=train(cls)
    if initial_reference is None: initial_reference=initial
    for a,b in zip(initial,initial_reference): torch.testing.assert_close(a,b,atol=0,rtol=0)
    assert np.isfinite(history).all()
    for index,ax in enumerate(axes): ax.plot(np.arange(1,81),history[:,index],label=cls.__name__)
    print(cls.__name__, "parameters:",count,"final train/validation MSE:",history[-1])
for ax,title in zip(axes,["Обучение","Отложенная выборка"] if ru else ["Training","Held-out validation"]):
    ax.set_title(title,loc="left");ax.set_yscale("log")
    ax.set_xlabel("Шаг" if ru else "Step");ax.set_ylabel("MSE");ax.grid(alpha=.25);ax.legend()
fig.tight_layout();plt.show();plt.close(fig)

# Paired scores isolate scale; they are not a complete attention matrix.
torch.manual_seed(SEED)
q = torch.randn(256, 8, 64) * 2.5
k = torch.randn(256, 8, 64) * 2.5
raw = (q*k).sum(-1)/math.sqrt(64)
q_n=F.rms_norm(q,[64],eps=1e-6);k_n=F.rms_norm(k,[64],eps=1e-6)
rms_scores=(q_n*k_n).sum(-1)/math.sqrt(64)
cosine=(F.normalize(q,dim=-1)*F.normalize(k,dim=-1)).sum(-1)
# A learned logit scale is intentionally omitted from this diagnostic.
torch.testing.assert_close(rms_scores,math.sqrt(64)*cosine,atol=2e-6,rtol=2e-6)
assert rms_scores.std() < raw.std()
print("raw / RMS / cosine std:",*[float(z.std()) for z in (raw,rms_scores,cosine)])
changed=q.clone();changed[:,1]*=10
per_head=F.rms_norm(changed,[64],eps=1e-6)
joint=F.rms_norm(q.flatten(-2),[8*64],eps=1e-6).reshape_as(q)
joint_changed=F.rms_norm(changed.flatten(-2),[8*64],eps=1e-6).reshape_as(q)
torch.testing.assert_close(per_head[:,0],q_n[:,0],atol=0,rtol=0)
assert not torch.allclose(joint[:,0],joint_changed[:,0])
print("Ось нормы меняет связь между головами." if ru else "The reduction axis changes coupling between heads.")