Модуль 8. Альтернативы Transformer: SSM, Mamba и гибриды

import hashlib
import json
import math
import subprocess
import sys
from dataclasses import dataclass
from numbers import Integral, Real
from typing import Any, Iterable

try:
    import numpy as np
except ImportError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "numpy>=1.24,<3"])
    import numpy as np

np.set_printoptions(precision=6, suppress=True)
REFERENCE_RAW = '{"affine_scan":{"chunk_size":16,"max_abs_difference":8.881784197001252e-16,"output_checksum":2.5500288962870172,"seed":11,"sequence_length":53,"state_width":6},"decode_flops":{"attention_core_flops":[16777216,536870912,17179869184],"inner_width":8192,"model_width":4096,"positions":[1024,32768,1048576],"ratios":[25.6,819.2,26214.4],"scope":"leading-order attention matmuls versus diagonal state update plus readout, with prepared coefficients; not end-to-end latency","ssm_core_flops":[655360,655360,655360],"state_width":16},"delta":{"additive_retrieval":[4.0,6.0],"full_retrieval":[3.0,4.0],"gated_half_retrieval":[1.9000000000000001,2.8],"half_retrieval":[2.0,3.0]},"linear_attention":{"chunk_size":16,"key_width":8,"max_abs_difference":2.1316282072803006e-14,"output_checksum":111.6891373299517,"seed":19,"sequence_length":50,"value_width":6},"memory":{"full_attention_256k_bytes":34359738368,"hybrid_256k":{"attention_layers":4,"conv_bytes":1835008,"kv_bytes":4294967296,"recurrent_bytes":9175040,"ssm_bytes":7340032,"ssm_layers":28,"total_bytes":4304142336},"kv_bytes_per_token_full_attention":131072,"mamba1":{"conv_bytes":2097152,"crossover_first_integer_token":80,"crossover_tokens_exact":80.0,"first_strictly_larger_token":81,"ssm_bytes":8388608,"total_bytes":10485760},"mamba2_like":{"conv_bytes":2162688,"crossover_first_integer_token":529,"crossover_tokens_exact":528.5,"first_strictly_larger_token":529,"ssm_bytes":67108864,"total_bytes":69271552}},"recall":{"expected_additive_mse":[0.09375,0.21875,0.46875,0.96875,1.96875,3.96875],"key_width":32,"mse":[0.13978625740853395,0.21200348269044997,0.5995392976261167,0.9892555844637765,1.8757385532834472,4.141977097591836],"paired_softmax_mse":[3.7460446307862835e-10,4.0983742994017256e-12,9.826283934707642e-10,1.7655760849531266e-09,1.1566029996764256e-09,1.0724543396718937e-08],"seed":7,"softmax_inverse_temperature":20.0,"softmax_mse_at_128":1.648958588765511e-08,"softmax_seed":107,"stored_pairs":[4,8,16,32,64,128],"value_width":32},"schema_version":2,"zoh":{"approximate_outputs":[0.2,-1.9496785275591946,1.0528816219675536],"approximate_states":[[0.1,0.1],[-0.9632120558828557,-0.9864664716763387],[0.4157823563022096,0.637099265665344]],"continuous_a":[-1.0,-2.0],"continuous_b":[[1.0],[1.0]],"deltas":[0.1,1.0,0.5],"exact_input_matrices":[[[0.09516258196404043],[0.09063462346100906]],[[0.6321205588285577],[0.43233235838169365]],[[0.39346934028736663],[0.31606027941427883]]],"exact_outputs":[0.1857972054250495,-1.0171784972997493,0.9023586277858214],"exact_states":[[0.09516258196404043,0.09063462346100906],[-0.5971122013551949,-0.42006629594455425],[0.42477182316430406,0.47758680462151737]],"inputs":[[1.0],[-1.0],[2.0]],"output_c":[[1.0,1.0]],"relative_input_error_first_mode":[0.05083319447750502,0.5819767068693265,0.270747041268399],"ten_percent_relative_error_root_a_minus_1":0.1937475579949906,"transitions":[[0.9048374180359595,0.8187307530779818],[0.36787944117144233,0.1353352832366127],[0.6065306597126334,0.36787944117144233]],"zero_mode_input_matrix":[[0.25],[0.22119921692859512]]}}'
EXPECTED_SHA256 = '28cbfd26cb5163eff23dd9ac6fd3084a325ee3a5efcb96192edd9562ae9b7387'
actual_sha256 = hashlib.sha256(REFERENCE_RAW.encode("utf-8")).hexdigest()
assert actual_sha256 == EXPECTED_SHA256
REFERENCE = json.loads(REFERENCE_RAW)
assert REFERENCE["schema_version"] == 2
print('Встроенный контракт М8 проверен', actual_sha256[:12])


"""Minimal M8 operators; used by the validated reference and embedded in notebooks."""
import numpy as np

def zoh_diagonal(a, b, delta):
    """Diagonal A: [n], B: [n,p], positive scalar step; caller validates shapes."""
    a = np.asarray(a, dtype=np.float64)
    b = np.asarray(b, dtype=np.float64)
    transition = np.exp(delta * a)
    factor = np.full_like(a, float(delta))
    np.divide(np.expm1(delta * a), a, out=factor, where=a != 0)
    return transition, factor[:, None] * b

def delta_step(state, key, value, beta, decay=1.0):
    """State [dv,dk], unit key [dk], value [dv]; decay before correction."""
    base = decay * state
    prediction = base @ key
    return base + beta * np.outer(value - prediction, key)

def compose_affine_pair(later, earlier):
    """Elementwise transforms; returns later(earlier(h)), not the reverse."""
    a2, b2 = later
    a1, b1 = earlier
    return a2 * a1, a2 * b1 + b2

def _real(name: str, value: Real) -> float:
    """Return one finite real scalar, rejecting booleans and non-finite values."""

    if isinstance(value, (bool, np.bool_)) or not isinstance(value, Real):
        raise TypeError(f"{name} must be a real scalar")
    number = float(value)
    if not math.isfinite(number):
        raise ValueError(f"{name} must be finite")
    return number


def _positive_real(name: str, value: Real) -> float:
    """Return one finite strictly positive scalar."""

    number = _real(name, value)
    if number <= 0.0:
        raise ValueError(f"{name} must be positive")
    return number


def _unit_interval(name: str, value: Real) -> float:
    """Return a finite scalar in the closed unit interval."""

    number = _real(name, value)
    if number < 0.0 or number > 1.0:
        raise ValueError(f"{name} must lie in [0, 1]")
    return number


def _positive_int(name: str, value: Integral) -> int:
    """Return one strictly positive integer, rejecting booleans."""

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


def _non_negative_int(name: str, value: Integral) -> int:
    """Return one non-negative integer, rejecting booleans."""

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


def _finite_array(
    name: str,
    value: Any,
    *,
    ndim: int | None = None,
    non_empty: bool = True,
) -> np.ndarray:
    """Return a float64 array with explicit dimensionality and finiteness checks."""

    array = np.asarray(value, dtype=np.float64)
    if ndim is not None and array.ndim != ndim:
        raise ValueError(f"{name} must have ndim={ndim}, got shape {array.shape}")
    if non_empty and array.size == 0:
        raise ValueError(f"{name} must be non-empty")
    if not np.all(np.isfinite(array)):
        raise ValueError(f"{name} must contain only finite values")
    return array


def bytes_to_kib(byte_count: Real) -> float:
    """Convert a finite non-negative byte count to binary KiB."""

    number = _real("byte_count", byte_count)
    if number < 0.0:
        raise ValueError("byte_count must be non-negative")
    return number / 2**10


def bytes_to_mib(byte_count: Real) -> float:
    """Convert a finite non-negative byte count to binary MiB."""

    number = _real("byte_count", byte_count)
    if number < 0.0:
        raise ValueError("byte_count must be non-negative")
    return number / 2**20


def bytes_to_gib(byte_count: Real) -> float:
    """Convert a finite non-negative byte count to binary GiB."""

    number = _real("byte_count", byte_count)
    if number < 0.0:
        raise ValueError("byte_count must be non-negative")
    return number / 2**30

@dataclass(frozen=True)
class ZOHDiscretization:
    """Exact diagonal zero-order-hold discretization for one time step.

    ``transition`` has shape ``[d_state]`` and represents the diagonal of
    ``exp(delta*A)``. ``input_matrix`` has shape ``[d_state, d_input]``.
    """

    transition: np.ndarray
    input_matrix: np.ndarray


@dataclass(frozen=True)
class SSMScanResult:
    """States, outputs, and per-step discrete parameters of a diagonal SSM scan."""

    states: np.ndarray
    outputs: np.ndarray
    transitions: np.ndarray
    input_matrices: np.ndarray


def zoh_discretize_diagonal(
    continuous_a: Any,
    continuous_b: Any,
    delta: Real,
) -> ZOHDiscretization:
    """Discretize ``dh/dt = A h + B x`` for diagonal stable ``A``.

    Parameters
    ----------
    continuous_a:
        Diagonal entries of ``A`` with shape ``[d_state]``. Non-positive values are
        accepted; positive values are rejected because this lesson models stable modes.
    continuous_b:
        Input matrix with shape ``[d_state, d_input]``. A vector is treated as one input
        channel.
    delta:
        Strictly positive time step.

    Notes
    -----
    The exact input factor is ``expm1(delta*a)/a``. At ``a=0`` the continuous limit is
    exactly ``delta``; handling that limit explicitly avoids a hidden ``0/0``.
    """

    a = _finite_array("continuous_a", continuous_a, ndim=1)
    if np.any(a > 0.0):
        raise ValueError("continuous_a must be non-positive for the stable lesson model")
    b = _finite_array("continuous_b", continuous_b)
    if b.ndim == 1:
        b = b[:, None]
    if b.ndim != 2 or b.shape[0] != a.size:
        raise ValueError("continuous_b must have shape [d_state, d_input]")
    step = _positive_real("delta", delta)

    transition, input_matrix = zoh_diagonal(a, b, step)
    if not np.all(np.isfinite(input_matrix)):
        raise FloatingPointError("ZOH discretization produced a non-finite input matrix")
    return ZOHDiscretization(transition=transition, input_matrix=input_matrix)


def scan_diagonal_ssm(
    continuous_a: Any,
    continuous_b: Any,
    output_c: Any,
    inputs: Any,
    deltas: Any,
    *,
    initial_state: Any | None = None,
    exact_input_discretization: bool = True,
) -> SSMScanResult:
    """Run a selective diagonal SSM over a non-empty sequence.

    Shapes are ``A:[d_state]``, ``B:[d_state,d_input]``, ``C:[d_output,d_state]``,
    ``inputs:[T,d_input]``, and ``deltas:[T]``. A one-dimensional input sequence is
    interpreted as ``d_input=1``. The output includes every post-update state.

    Setting ``exact_input_discretization=False`` keeps the exact transition but replaces
    the ZOH input matrix by the common small-step approximation ``delta*B``.
    """

    a = _finite_array("continuous_a", continuous_a, ndim=1)
    b = _finite_array("continuous_b", continuous_b)
    if b.ndim == 1:
        b = b[:, None]
    c = _finite_array("output_c", output_c)
    if c.ndim == 1:
        c = c[None, :]
    x = _finite_array("inputs", inputs)
    if x.ndim == 1:
        x = x[:, None]
    dt = _finite_array("deltas", deltas, ndim=1)

    if b.ndim != 2 or b.shape[0] != a.size:
        raise ValueError("continuous_b must have shape [d_state, d_input]")
    if c.ndim != 2 or c.shape[1] != a.size:
        raise ValueError("output_c must have shape [d_output, d_state]")
    if x.ndim != 2 or x.shape[1] != b.shape[1]:
        raise ValueError("inputs must have shape [T, d_input]")
    if x.shape[0] != dt.size:
        raise ValueError("inputs and deltas must have the same sequence length")
    if np.any(dt <= 0.0):
        raise ValueError("every delta must be positive")
    if np.any(a > 0.0):
        raise ValueError("continuous_a must be non-positive for the stable lesson model")

    state = (
        np.zeros(a.size, dtype=np.float64)
        if initial_state is None
        else _finite_array("initial_state", initial_state, ndim=1).copy()
    )
    if state.shape != a.shape:
        raise ValueError("initial_state must have shape [d_state]")

    states = np.empty((x.shape[0], a.size), dtype=np.float64)
    outputs = np.empty((x.shape[0], c.shape[0]), dtype=np.float64)
    transitions = np.empty_like(states)
    input_matrices = np.empty((x.shape[0], a.size, b.shape[1]), dtype=np.float64)

    for index, (token, step) in enumerate(zip(x, dt)):
        discrete = zoh_discretize_diagonal(a, b, float(step))
        input_matrix = (
            discrete.input_matrix
            if exact_input_discretization
            else float(step) * b
        )
        state = discrete.transition * state + input_matrix @ token
        states[index] = state
        outputs[index] = c @ state
        transitions[index] = discrete.transition
        input_matrices[index] = input_matrix

    if not np.all(np.isfinite(states)) or not np.all(np.isfinite(outputs)):
        raise FloatingPointError("SSM scan produced non-finite states or outputs")
    return SSMScanResult(states, outputs, transitions, input_matrices)


z = REFERENCE["zoh"]
a = np.asarray(z["continuous_a"], dtype=np.float64)
b = np.asarray(z["continuous_b"], dtype=np.float64)
c = np.asarray(z["output_c"], dtype=np.float64)
x = np.asarray(z["inputs"], dtype=np.float64)
delta = np.asarray(z["deltas"], dtype=np.float64)

exact = scan_diagonal_ssm(a, b, c, x, delta)
approx = scan_diagonal_ssm(a, b, c, x, delta, exact_input_discretization=False)
assert exact.states.shape == (3, 2)
assert exact.outputs.shape == (3, 1)
assert np.allclose(exact.transitions, z["transitions"], rtol=1e-12, atol=1e-12)
assert np.allclose(exact.input_matrices, z["exact_input_matrices"], rtol=1e-12, atol=1e-12)
assert np.allclose(exact.states, z["exact_states"], rtol=1e-12, atol=1e-12)
assert np.allclose(exact.outputs[:, 0], z["exact_outputs"], rtol=1e-12, atol=1e-12)
assert np.allclose(approx.states, z["approximate_states"], rtol=1e-12, atol=1e-12)

relative_errors = []
for step in delta:
    discrete = zoh_discretize_diagonal(a, b, float(step))
    relative_errors.append(abs(step * b[0, 0] - discrete.input_matrix[0, 0]) / discrete.input_matrix[0, 0])
assert np.allclose(relative_errors, z["relative_input_error_first_mode"], rtol=1e-12, atol=1e-12)

# The zero continuous mode uses the analytical limit phi(0)=delta.
zero_mode = zoh_discretize_diagonal([0.0, -1.0], [[1.0], [1.0]], 0.25)
assert np.allclose(zero_mode.input_matrix, z["zero_mode_input_matrix"], atol=1e-12)
assert zero_mode.input_matrix[0, 0] == 0.25
# A small a is not a zero mode when delta*a is order one.
with np.errstate(divide="raise", invalid="raise"):
    long_step = zoh_discretize_diagonal([-1e-15, 0.0], [[1.0], [1.0]], 1e15)
np.testing.assert_allclose(long_step.input_matrix[:, 0], [(1-np.exp(-1))*1e15, 1e15], rtol=1e-14)


root = z["ten_percent_relative_error_root_a_minus_1"]
root_error = root / (-math.expm1(-root)) - 1.0
assert math.isclose(root, 0.1937475579949906, rel_tol=0.0, abs_tol=1e-13)
assert math.isclose(root_error, 0.1, rel_tol=0.0, abs_tol=1e-12)

for index, step in enumerate(delta):
    print(
        f"t={index+1} delta={step:.1f} | Abar={exact.transitions[index]} | "
        f"relative ΔB error={relative_errors[index]:.1%} | y={exact.outputs[index,0]:+.6f}"
    )
print('ZOH, селективный проход и предел при a=0 проверены')

@dataclass(frozen=True)
class AffineTransform:
    """Elementwise affine state transform ``h_out = multiplier*h_in + bias``."""

    multiplier: np.ndarray
    bias: np.ndarray


def compose_affine(later: AffineTransform, earlier: AffineTransform) -> AffineTransform:
    """Compose two diagonal affine transforms as ``later(earlier(h))``."""

    a2 = _finite_array("later.multiplier", later.multiplier, ndim=1)
    b2 = _finite_array("later.bias", later.bias, ndim=1)
    a1 = _finite_array("earlier.multiplier", earlier.multiplier, ndim=1)
    b1 = _finite_array("earlier.bias", earlier.bias, ndim=1)
    if a1.shape != b1.shape or a2.shape != b2.shape or a1.shape != a2.shape:
        raise ValueError("all affine components must share shape [d_state]")
    return AffineTransform(multiplier=a2 * a1, bias=a2 * b1 + b2)


def affine_scan_sequential(
    multipliers: Any,
    biases: Any,
    *,
    initial_state: Any | None = None,
) -> np.ndarray:
    """Evaluate ``h_t = a_t*h_{t-1} + b_t`` sequentially.

    ``multipliers`` and ``biases`` must both have shape ``[T, d_state]``. The returned
    array contains post-update states with the same shape.
    """

    a = _finite_array("multipliers", multipliers, ndim=2)
    b = _finite_array("biases", biases, ndim=2)
    if a.shape != b.shape:
        raise ValueError("multipliers and biases must have identical shape")
    state = (
        np.zeros(a.shape[1], dtype=np.float64)
        if initial_state is None
        else _finite_array("initial_state", initial_state, ndim=1).copy()
    )
    if state.shape != (a.shape[1],):
        raise ValueError("initial_state must have shape [d_state]")
    result = np.empty_like(b)
    for index in range(a.shape[0]):
        state = a[index] * state + b[index]
        result[index] = state
    return result


def affine_scan_chunked(
    multipliers: Any,
    biases: Any,
    *,
    chunk_size: Integral,
    initial_state: Any | None = None,
) -> np.ndarray:
    """Evaluate the same affine recurrence in non-overlapping chunks.

    Inside each chunk, the local zero-state recurrence and cumulative multipliers can be
    computed in parallel by a scan kernel. Only one boundary state crosses each chunk.
    This NumPy version keeps the mechanics transparent; it is not a fused GPU scan.
    """

    a = _finite_array("multipliers", multipliers, ndim=2)
    b = _finite_array("biases", biases, ndim=2)
    if a.shape != b.shape:
        raise ValueError("multipliers and biases must have identical shape")
    chunk = _positive_int("chunk_size", chunk_size)
    state = (
        np.zeros(a.shape[1], dtype=np.float64)
        if initial_state is None
        else _finite_array("initial_state", initial_state, ndim=1).copy()
    )
    if state.shape != (a.shape[1],):
        raise ValueError("initial_state must have shape [d_state]")

    result = np.empty_like(b)
    for start in range(0, a.shape[0], chunk):
        end = min(start + chunk, a.shape[0])
        local_a = a[start:end]
        local_b = b[start:end]

        # Zero-start local states carry all within-chunk contributions.
        zero_state = np.zeros(a.shape[1], dtype=np.float64)
        local_zero = np.empty_like(local_b)
        for offset in range(local_a.shape[0]):
            zero_state = local_a[offset] * zero_state + local_b[offset]
            local_zero[offset] = zero_state

        # The chunk-start state is multiplied by the cumulative transition to each row.
        prefix_multiplier = np.cumprod(local_a, axis=0)
        result[start:end] = local_zero + prefix_multiplier * state
        state = result[end - 1].copy()
    return result


ref = REFERENCE["affine_scan"]
rng = np.random.default_rng(ref["seed"])
a_steps = rng.uniform(0.70, 0.99, size=(ref["sequence_length"], ref["state_width"]))
b_steps = rng.standard_normal((ref["sequence_length"], ref["state_width"]))
h0 = rng.standard_normal(ref["state_width"])
seq = affine_scan_sequential(a_steps, b_steps, initial_state=h0)
chunked = affine_scan_chunked(a_steps, b_steps, chunk_size=ref["chunk_size"], initial_state=h0)
difference = float(np.max(np.abs(seq - chunked)))
assert difference < 1e-12
assert math.isclose(float(np.sum(seq)), ref["output_checksum"], rel_tol=1e-11, abs_tol=1e-12)

# Composition itself is associative; check it on three transforms.
T1 = AffineTransform(np.array([0.9, 0.8]), np.array([1.0, -1.0]))
T2 = AffineTransform(np.array([0.7, 0.6]), np.array([0.5, 2.0]))
T3 = AffineTransform(np.array([0.5, 0.4]), np.array([-0.5, 1.0]))
left = compose_affine(T3, compose_affine(T2, T1))
right = compose_affine(compose_affine(T3, T2), T1)
assert np.allclose(left.multiplier, right.multiplier, atol=1e-15)
assert np.allclose(left.bias, right.bias, atol=1e-15)
print(f"max|sequential - chunked| = {difference:.3e}; uneven tail = {ref['sequence_length'] % ref['chunk_size']}")
print('Affine associativity и блочный проход проверены')

def _validate_qkv(q: Any, k: Any, v: Any) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Validate a non-empty causal linear-attention triplet."""

    query = _finite_array("q", q, ndim=2)
    key = _finite_array("k", k, ndim=2)
    value = _finite_array("v", v, ndim=2)
    if query.shape[0] != key.shape[0] or query.shape[0] != value.shape[0]:
        raise ValueError("q, k, and v must share sequence length T")
    if query.shape[1] != key.shape[1]:
        raise ValueError("q and k must share feature width d_key")
    return query, key, value


def linear_attention_sequential(q: Any, k: Any, v: Any) -> np.ndarray:
    """Causal additive linear attention with state shape ``[d_value, d_key]``."""

    query, key, value = _validate_qkv(q, k, v)
    state = np.zeros((value.shape[1], key.shape[1]), dtype=np.float64)
    output = np.empty((query.shape[0], value.shape[1]), dtype=np.float64)
    for index in range(query.shape[0]):
        state += np.outer(value[index], key[index])
        output[index] = state @ query[index]
    return output


def linear_attention_chunked(
    q: Any,
    k: Any,
    v: Any,
    *,
    chunk_size: Integral,
) -> np.ndarray:
    """Exact chunked form of the same causal additive linear attention.

    The inter-chunk term reads the carried state; the intra-chunk term is a masked
    matrix multiplication. The final partial chunk is supported explicitly.
    """

    query, key, value = _validate_qkv(q, k, v)
    chunk = _positive_int("chunk_size", chunk_size)
    state = np.zeros((value.shape[1], key.shape[1]), dtype=np.float64)
    output = np.empty((query.shape[0], value.shape[1]), dtype=np.float64)
    for start in range(0, query.shape[0], chunk):
        end = min(start + chunk, query.shape[0])
        q_chunk = query[start:end]
        k_chunk = key[start:end]
        v_chunk = value[start:end]
        inter = q_chunk @ state.T
        causal_scores = np.tril(q_chunk @ k_chunk.T)
        output[start:end] = inter + causal_scores @ v_chunk
        state += v_chunk.T @ k_chunk
    return output


ref = REFERENCE["linear_attention"]
rng = np.random.default_rng(ref["seed"])
q = rng.standard_normal((ref["sequence_length"], ref["key_width"]))
k = rng.standard_normal((ref["sequence_length"], ref["key_width"]))
v = rng.standard_normal((ref["sequence_length"], ref["value_width"]))
sequential = linear_attention_sequential(q, k, v)
for chunk_size in (1, 7, ref["chunk_size"], 64):
    candidate = linear_attention_chunked(q, k, v, chunk_size=chunk_size)
    assert candidate.shape == sequential.shape
    assert np.allclose(candidate, sequential, rtol=1e-12, atol=1e-12)
max_difference = float(np.max(np.abs(linear_attention_chunked(q, k, v, chunk_size=16) - sequential)))
assert max_difference < 1e-12
assert math.isclose(float(np.sum(sequential)), ref["output_checksum"], rel_tol=1e-11, abs_tol=1e-12)
# Perturbing future tokens must not change earlier causal outputs.
changed = v.copy(); changed[-1] += 100
np.testing.assert_allclose(linear_attention_chunked(q, k, changed, chunk_size=7)[:-1], sequential[:-1])
print(f"output shape={sequential.shape}, max difference={max_difference:.3e}")
print('Sequential/chunked linear attention совпадают')

@dataclass(frozen=True)
class DeltaUpdateResult:
    """State and retrieval trace for one normalized-key delta-rule update."""

    state: np.ndarray
    normalized_key: np.ndarray
    retrieval_before_write: np.ndarray
    retrieval_after_write: np.ndarray


def normalize_key(key: Any) -> np.ndarray:
    """Return a unit-L2 key, rejecting the zero vector."""

    vector = _finite_array("key", key, ndim=1)
    scale = float(np.max(np.abs(vector)))
    if scale == 0.0:
        raise ValueError("key must have non-zero norm")
    scaled = vector / scale
    return scaled / np.linalg.norm(scaled)


def delta_update(
    state: Any,
    key: Any,
    value: Any,
    *,
    beta: Real,
    decay: Real = 1.0,
    normalize: bool = True,
) -> DeltaUpdateResult:
    """Apply a gated delta-rule write to associative state ``[d_value, d_key]``.

    With a unit key, the update is

    ``base = decay * state``
    ``new = base + beta * outer(value - base@key, key)``.

    Thus ``beta=1`` overwrites retrieval at that key exactly, while ``beta=0`` only
    applies the global decay. If ``normalize=False``, the caller must supply a unit key;
    the explicit check prevents a misleading "exact overwrite" claim for arbitrary norm.
    """

    matrix = _finite_array("state", state, ndim=2)
    raw_key = _finite_array("key", key, ndim=1)
    target = _finite_array("value", value, ndim=1)
    if matrix.shape != (target.size, raw_key.size):
        raise ValueError("state must have shape [d_value, d_key]")
    write_rate = _unit_interval("beta", beta)
    forget_rate = _unit_interval("decay", decay)
    unit_key = normalize_key(raw_key) if normalize else raw_key.copy()
    if not normalize and not np.isclose(np.linalg.norm(unit_key), 1.0, rtol=0.0, atol=1e-10):
        raise ValueError("key must have unit norm when normalize=False")

    base = forget_rate * matrix
    before = base @ unit_key
    updated = delta_step(matrix, unit_key, target, write_rate, forget_rate)
    after = updated @ unit_key
    return DeltaUpdateResult(updated, unit_key, before, after)


ref = REFERENCE["delta"]
k = np.array([1.0, 0.0])
v1 = np.array([1.0, 2.0])
v2 = np.array([3.0, 4.0])
state = np.outer(v1, k)
additive = state + np.outer(v2, k)
half = delta_update(state, k, v2, beta=0.5)
full = delta_update(state, k, v2, beta=1.0)
gated = delta_update(state, k, v2, beta=0.5, decay=0.8)
no_write = delta_update(state, k, v2, beta=0.0, decay=0.5)

assert np.allclose(additive @ k, ref["additive_retrieval"])
assert np.allclose(half.retrieval_after_write, ref["half_retrieval"])
assert np.allclose(full.retrieval_after_write, ref["full_retrieval"])
assert np.allclose(gated.retrieval_after_write, ref["gated_half_retrieval"])
assert np.allclose(full.retrieval_after_write, v2)  # exact overwrite for a unit key
assert np.allclose(no_write.state, 0.5 * state)

# normalize=True makes a scaled copy of the same key semantically identical.
scaled_key = delta_update(state, 7.0 * k, v2, beta=1.0, normalize=True)
assert np.allclose(scaled_key.state, full.state)
for magnitude in (1e300, 1e-300):
    np.testing.assert_allclose(normalize_key([magnitude, magnitude]), np.ones(2)/np.sqrt(2))
# Targeted overwrite can change another, non-orthogonal key.
other = np.array([1., 1.]) / np.sqrt(2)
assert not np.allclose(full.state @ other, state @ other)

print("additive", additive @ k, "| beta=.5", half.retrieval_after_write, "| beta=1", full.retrieval_after_write)
print('Delta rule, decay gate и нормировка ключа проверены')

@dataclass(frozen=True)
class RecallCurve:
    """Point-retrieval MSE for additive associative memory at several loads."""

    stored_pairs: tuple[int, ...]
    mse: tuple[float, ...]
    key_width: int
    value_width: int
    seed: int


def associative_recall_curve(
    stored_pairs: Iterable[Integral],
    *,
    key_width: Integral,
    value_width: Integral | None = None,
    seed: Integral = 7,
) -> RecallCurve:
    """Measure interference in ``S = sum(v k^T)`` for normalized random keys."""

    widths = _positive_int("key_width", key_width)
    values = widths if value_width is None else _positive_int("value_width", value_width)
    random_seed = _non_negative_int("seed", seed)
    loads = tuple(_positive_int("stored_pairs item", item) for item in stored_pairs)
    if not loads:
        raise ValueError("stored_pairs must be non-empty")

    rng = np.random.default_rng(random_seed)
    errors: list[float] = []
    for count in loads:
        keys = rng.standard_normal((count, widths))
        keys /= np.linalg.norm(keys, axis=1, keepdims=True)
        targets = rng.standard_normal((count, values))
        state = targets.T @ keys
        retrieved = (state @ keys.T).T
        errors.append(float(np.mean((retrieved - targets) ** 2)))
    return RecallCurve(loads, tuple(errors), widths, values, random_seed)


def stable_softmax_rows(logits: Any) -> np.ndarray:
    """Compute row-wise softmax for a finite non-empty matrix."""

    scores = _finite_array("logits", logits, ndim=2)
    shifted = scores - np.max(scores, axis=1, keepdims=True)
    weights = np.exp(shifted)
    denominator = np.sum(weights, axis=1, keepdims=True)
    if np.any(denominator <= 0.0) or not np.all(np.isfinite(denominator)):
        raise FloatingPointError("softmax denominator must be positive and finite")
    return weights / denominator


def softmax_addressing_mse(
    *,
    stored_pairs: Integral,
    key_width: Integral,
    value_width: Integral | None = None,
    inverse_temperature: Real = 20.0,
    seed: Integral = 107,
) -> float:
    """Return finite-temperature self-retrieval MSE from an explicit key/value table.

    This is a near-one-hot addressing contrast, not a claim that finite-temperature
    softmax is mathematically exact for arbitrary random keys.
    """

    count = _positive_int("stored_pairs", stored_pairs)
    width = _positive_int("key_width", key_width)
    values = width if value_width is None else _positive_int("value_width", value_width)
    temperature = _positive_real("inverse_temperature", inverse_temperature)
    random_seed = _non_negative_int("seed", seed)
    rng = np.random.default_rng(random_seed)
    keys = rng.standard_normal((count, width))
    keys /= np.linalg.norm(keys, axis=1, keepdims=True)
    targets = rng.standard_normal((count, values))
    weights = stable_softmax_rows(temperature * (keys @ keys.T))
    return float(np.mean((weights @ targets - targets) ** 2))


ref = REFERENCE["recall"]
curve = associative_recall_curve(
    ref["stored_pairs"],
    key_width=ref["key_width"],
    value_width=ref["value_width"],
    seed=ref["seed"],
)
assert np.allclose(curve.mse, ref["mse"], rtol=1e-11, atol=1e-12)
assert all(later > earlier for earlier, later in zip(curve.mse[2:], curve.mse[3:]))
softmax_mse = softmax_addressing_mse(
    stored_pairs=128,
    key_width=32,
    value_width=32,
    inverse_temperature=ref["softmax_inverse_temperature"],
    seed=ref["softmax_seed"],
)
assert math.isclose(softmax_mse, ref["softmax_mse_at_128"], rel_tol=1e-10, abs_tol=1e-12)
assert 0.0 < softmax_mse < 1e-6
# The plotted control uses the same keys and values at every load.
paired_rng = np.random.default_rng(ref["seed"])
paired_errors = []
for count in ref["stored_pairs"]:
    keys = paired_rng.standard_normal((count, 32))
    keys /= np.linalg.norm(keys, axis=1, keepdims=True)
    values = paired_rng.standard_normal((count, 32))
    prediction = stable_softmax_rows(20 * keys @ keys.T) @ values
    paired_errors.append(float(np.mean((prediction-values)**2)))
np.testing.assert_allclose(paired_errors, ref["paired_softmax_mse"], rtol=1e-10, atol=1e-15)
print("linear-state MSE:", [round(value, 4) for value in curve.mse])
print(f"finite-temperature softmax MSE at 128 pairs: {softmax_mse:.3e}")
print('Recall-интерференция и finite-temperature softmax проверены')

@dataclass(frozen=True)
class SSMStateGeometry:
    """Per-sequence inference-cache geometry for an SSM stack.

    A recurrent Mamba-style layer can retain both the main SSM state and the
    short causal-convolution state. ``conv_channels`` is explicit because
    Mamba-2 may cache x plus B/C channels rather than only ``inner_width``.
    """

    layers: int
    inner_width: int
    state_width: int
    conv_width: int = 0
    conv_channels: int = 0
    bytes_per_value: int = 2

    def __post_init__(self) -> None:
        _positive_int("layers", self.layers)
        _positive_int("inner_width", self.inner_width)
        _positive_int("state_width", self.state_width)
        _positive_int("bytes_per_value", self.bytes_per_value)
        _non_negative_int("conv_width", self.conv_width)
        _non_negative_int("conv_channels", self.conv_channels)
        if (self.conv_width == 0) != (self.conv_channels == 0):
            raise ValueError("conv_width and conv_channels must both be zero or both be positive")

    def ssm_bytes_per_sequence(self, *, batch_size: Integral = 1) -> int:
        """Return the main recurrent SSM-state payload."""

        batch = _positive_int("batch_size", batch_size)
        return batch * self.layers * self.inner_width * self.state_width * self.bytes_per_value

    def conv_bytes_per_sequence(self, *, batch_size: Integral = 1) -> int:
        """Return the short causal-convolution cache payload."""

        batch = _positive_int("batch_size", batch_size)
        return batch * self.layers * self.conv_channels * self.conv_width * self.bytes_per_value

    def bytes_per_sequence(self, *, batch_size: Integral = 1) -> int:
        """Return the complete fixed recurrent cache for one or more sequences."""

        return self.ssm_bytes_per_sequence(batch_size=batch_size) + self.conv_bytes_per_sequence(batch_size=batch_size)


@dataclass(frozen=True)
class KVCacheGeometry:
    """Per-token K/V geometry for a stack of attention layers."""

    layers: int
    kv_heads: int
    head_width: int
    bytes_per_value: int = 2

    def __post_init__(self) -> None:
        _positive_int("layers", self.layers)
        _positive_int("kv_heads", self.kv_heads)
        _positive_int("head_width", self.head_width)
        _positive_int("bytes_per_value", self.bytes_per_value)

    @property
    def bytes_per_token(self) -> int:
        return 2 * self.layers * self.kv_heads * self.head_width * self.bytes_per_value

    def total_bytes(self, context_tokens: Integral, *, batch_size: Integral = 1) -> int:
        tokens = _non_negative_int("context_tokens", context_tokens)
        batch = _positive_int("batch_size", batch_size)
        return batch * tokens * self.bytes_per_token


@dataclass(frozen=True)
class HybridMemoryEstimate:
    """Per-sequence KV and recurrent cache of a hybrid stack."""

    total_layers: int
    attention_layers: int
    ssm_layers: int
    context_tokens: int
    kv_bytes: int
    ssm_bytes: int
    conv_bytes: int
    recurrent_bytes: int
    total_bytes: int


def hybrid_memory_estimate(
    *,
    total_layers: Integral,
    attention_layers: Integral,
    context_tokens: Integral,
    inner_width: Integral,
    state_width: Integral,
    conv_width: Integral,
    conv_channels: Integral,
    kv_heads: Integral,
    head_width: Integral,
    bytes_per_value: Integral = 2,
) -> HybridMemoryEstimate:
    """Return KV plus complete recurrent-cache payload for one hybrid sequence."""

    total = _positive_int("total_layers", total_layers)
    if isinstance(attention_layers, (bool, np.bool_)) or not isinstance(attention_layers, Integral):
        raise TypeError("attention_layers must be an integer")
    attention = int(attention_layers)
    if attention < 0 or attention > total:
        raise ValueError("attention_layers must lie in [0, total_layers]")
    tokens = _non_negative_int("context_tokens", context_tokens)
    inner = _positive_int("inner_width", inner_width)
    state = _positive_int("state_width", state_width)
    conv = _non_negative_int("conv_width", conv_width)
    channels = _non_negative_int("conv_channels", conv_channels)
    if (conv == 0) != (channels == 0):
        raise ValueError("conv_width and conv_channels must both be zero or both be positive")
    heads = _positive_int("kv_heads", kv_heads)
    width = _positive_int("head_width", head_width)
    item_bytes = _positive_int("bytes_per_value", bytes_per_value)
    ssm_layers = total - attention

    kv_bytes = 0 if attention == 0 else KVCacheGeometry(attention, heads, width, item_bytes).total_bytes(tokens)
    if ssm_layers == 0:
        ssm_bytes = conv_bytes = recurrent_bytes = 0
    else:
        geometry = SSMStateGeometry(
            ssm_layers, inner, state, conv, channels, item_bytes
        )
        ssm_bytes = geometry.ssm_bytes_per_sequence()
        conv_bytes = geometry.conv_bytes_per_sequence()
        recurrent_bytes = geometry.bytes_per_sequence()
    return HybridMemoryEstimate(
        total_layers=total,
        attention_layers=attention,
        ssm_layers=ssm_layers,
        context_tokens=tokens,
        kv_bytes=kv_bytes,
        ssm_bytes=ssm_bytes,
        conv_bytes=conv_bytes,
        recurrent_bytes=recurrent_bytes,
        total_bytes=kv_bytes + recurrent_bytes,
    )


@dataclass(frozen=True)
class DecodeFlopEstimate:
    """Arithmetic of the central token-mixing operators only.

    The numbers exclude projections, gating, convolution, FFN, weight traffic,
    communication and kernel efficiency. They must not be read as latency or
    end-to-end model speedups.
    """

    context_tokens: int
    attention_core_flops: int
    ssm_core_flops: int
    ratio: float


def decode_flop_estimate(
    *,
    context_tokens: Integral,
    model_width: Integral,
    inner_width: Integral,
    state_width: Integral,
) -> DecodeFlopEstimate:
    """Return the declared core-only scenario ``4*N*d`` versus ``5*d_inner*d_state``."""

    tokens = _positive_int("context_tokens", context_tokens)
    model = _positive_int("model_width", model_width)
    inner = _positive_int("inner_width", inner_width)
    state = _positive_int("state_width", state_width)
    attention = 4 * tokens * model
    # State update: two multiplies + one add; readout: multiply + add.
    # Leading-order counts omit one readout addition per channel.
    ssm = 5 * inner * state
    return DecodeFlopEstimate(tokens, attention, ssm, attention / ssm)


ref = REFERENCE["memory"]
kv = KVCacheGeometry(layers=32, kv_heads=8, head_width=128, bytes_per_value=2)
mamba1 = SSMStateGeometry(
    layers=32, inner_width=8192, state_width=16,
    conv_width=4, conv_channels=8192, bytes_per_value=2,
)
mamba2 = SSMStateGeometry(
    layers=32, inner_width=8192, state_width=128,
    conv_width=4, conv_channels=8192 + 2 * 128, bytes_per_value=2,
)
assert kv.bytes_per_token == ref["kv_bytes_per_token_full_attention"] == 128 * 2**10
assert mamba1.ssm_bytes_per_sequence() == ref["mamba1"]["ssm_bytes"] == 8 * 2**20
assert mamba1.conv_bytes_per_sequence() == ref["mamba1"]["conv_bytes"] == 2 * 2**20
assert mamba1.bytes_per_sequence() == ref["mamba1"]["total_bytes"] == 10 * 2**20
assert ref["mamba1"]["crossover_tokens_exact"] == 80
assert ref["mamba1"]["crossover_first_integer_token"] == 80  # first >=
assert ref["mamba1"]["first_strictly_larger_token"] == 81
assert kv.total_bytes(80) == mamba1.bytes_per_sequence()
assert kv.total_bytes(81) > mamba1.bytes_per_sequence()
assert mamba2.ssm_bytes_per_sequence() == ref["mamba2_like"]["ssm_bytes"] == 64 * 2**20
assert math.isclose(bytes_to_mib(mamba2.conv_bytes_per_sequence()), 2.0625)
assert mamba2.bytes_per_sequence() == ref["mamba2_like"]["total_bytes"]
assert ref["mamba2_like"]["crossover_tokens_exact"] == 528.5
assert ref["mamba2_like"]["crossover_first_integer_token"] == 529

full_256k = kv.total_bytes(256 * 1024)
hybrid = hybrid_memory_estimate(
    total_layers=32, attention_layers=4, context_tokens=256 * 1024,
    inner_width=8192, state_width=16, conv_width=4, conv_channels=8192,
    kv_heads=8, head_width=128, bytes_per_value=2,
)
assert full_256k == ref["full_attention_256k_bytes"] == 32 * 2**30
assert hybrid.kv_bytes == ref["hybrid_256k"]["kv_bytes"] == 4 * 2**30
assert hybrid.ssm_bytes == ref["hybrid_256k"]["ssm_bytes"] == 7 * 2**20
assert hybrid.conv_bytes == ref["hybrid_256k"]["conv_bytes"] == int(1.75 * 2**20)
assert hybrid.recurrent_bytes == ref["hybrid_256k"]["recurrent_bytes"] == int(8.75 * 2**20)
assert hybrid.total_bytes == ref["hybrid_256k"]["total_bytes"]

assert mamba1.bytes_per_sequence(batch_size=3) == 3 * mamba1.bytes_per_sequence()
assert kv.total_bytes(8192, batch_size=3) == 3 * kv.total_bytes(8192)

flop_ref = REFERENCE["decode_flops"]
flop_rows = [decode_flop_estimate(context_tokens=position, model_width=4096, inner_width=8192, state_width=16) for position in flop_ref["positions"]]
assert [row.attention_core_flops for row in flop_rows] == flop_ref["attention_core_flops"]
assert [row.ssm_core_flops for row in flop_rows] == flop_ref["ssm_core_flops"]
assert np.allclose([row.ratio for row in flop_rows], flop_ref["ratios"], rtol=1e-12)
assert "not end-to-end latency" in flop_ref["scope"]
# Independent arithmetic: two attention dot-product stages; SSM write + read.
for position,row in zip(flop_ref["positions"],flop_rows):
    assert row.attention_core_flops == 2 * (2 * position * 4096)
    update = 8192 * 16 * (2 + 1)
    readout = 8192 * (16 + 16)  # leading estimate of the reduction
    assert row.ssm_core_flops == update + readout
assert np.allclose([row.ratio for row in flop_rows], [25.6,819.2,26214.4])

print(f"KV/token={bytes_to_kib(kv.bytes_per_token):.0f} KiB")
print(f"Mamba-1 cache={bytes_to_mib(mamba1.bytes_per_sequence()):.0f} MiB; crossover=80")
print(f"Mamba-2-like cache={bytes_to_mib(mamba2.bytes_per_sequence()):.4f} MiB; crossover=528.5 / first integer 529")
print(f"256K hybrid={bytes_to_gib(hybrid.total_bytes):.6f} GiB (KV 4 GiB + recurrent 8.75 MiB)")
print("core-only ratios:", [round(row.ratio) for row in flop_rows])
print('SSM/KV/hybrid память и FLOPs-сценарий проверены')

def expect_failure(exception_type, function, *args, **kwargs):
    """Assert that one invalid educational input fails with the declared exception."""

    try:
        function(*args, **kwargs)
    except exception_type:
        return
    raise AssertionError(f"expected {exception_type.__name__} from {function.__name__}")

expect_failure(ValueError, zoh_discretize_diagonal, [0.1], [[1.0]], 0.1)
expect_failure(ValueError, zoh_discretize_diagonal, [-1.0], [[1.0]], 0.0)
expect_failure(ValueError, scan_diagonal_ssm, [-1.0], [[1.0]], [[1.0]], [], [])
expect_failure(ValueError, scan_diagonal_ssm, [-1.0], [[1.0]], [[1.0]], [[1.0], [2.0]], [0.1])
expect_failure(ValueError, affine_scan_chunked, [[0.9]], [[1.0]], chunk_size=0)
expect_failure(ValueError, linear_attention_sequential, [[1.0]], [[1.0, 2.0]], [[1.0]])
expect_failure(ValueError, normalize_key, [0.0, 0.0])
expect_failure(ValueError, delta_update, np.zeros((2, 2)), [2.0, 0.0], [1.0, 2.0], beta=1.0, normalize=False)
expect_failure(ValueError, stable_softmax_rows, [[0.0, np.nan]])
expect_failure(ValueError, hybrid_memory_estimate, total_layers=8, attention_layers=9, context_tokens=128, inner_width=64, state_width=16, conv_width=4, conv_channels=64, kv_heads=2, head_width=8)
expect_failure(TypeError, KVCacheGeometry, layers=True, kv_heads=8, head_width=128)
print('Негативные тесты пройдены')

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"воспроизводимый тестовый запуск; seed={SEED}")
DEMO_LANGUAGE = 'ru'

plt.rcParams.update({"font.family":"DejaVu Sans", "font.size":12, "axes.titlesize":15, "axes.labelsize":13, "xtick.labelsize":11, "ytick.labelsize":11, "figure.dpi":150})

import math
from dataclasses import dataclass
from numbers import Integral, Real
from typing import Any


def _real(name: str, value: Real) -> float:
    """Return one finite real scalar, rejecting booleans and non-finite values."""

    if isinstance(value, (bool, np.bool_)) or not isinstance(value, Real):
        raise TypeError(f"{name} must be a real scalar")
    number = float(value)
    if not math.isfinite(number):
        raise ValueError(f"{name} must be finite")
    return number


def _positive_real(name: str, value: Real) -> float:
    """Return one finite strictly positive scalar."""

    number = _real(name, value)
    if number <= 0.0:
        raise ValueError(f"{name} must be positive")
    return number


def _unit_interval(name: str, value: Real) -> float:
    """Return a finite scalar in the closed unit interval."""

    number = _real(name, value)
    if number < 0.0 or number > 1.0:
        raise ValueError(f"{name} must lie in [0, 1]")
    return number


def _positive_int(name: str, value: Integral) -> int:
    """Return one strictly positive integer, rejecting booleans."""

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


def _non_negative_int(name: str, value: Integral) -> int:
    """Return one non-negative integer, rejecting booleans."""

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


def _finite_array(
    name: str,
    value: Any,
    *,
    ndim: int | None = None,
    non_empty: bool = True,
) -> np.ndarray:
    """Return a float64 array with explicit dimensionality and finiteness checks."""

    array = np.asarray(value, dtype=np.float64)
    if ndim is not None and array.ndim != ndim:
        raise ValueError(f"{name} must have ndim={ndim}, got shape {array.shape}")
    if non_empty and array.size == 0:
        raise ValueError(f"{name} must be non-empty")
    if not np.all(np.isfinite(array)):
        raise ValueError(f"{name} must contain only finite values")
    return array


def bytes_to_kib(byte_count: Real) -> float:
    """Convert a finite non-negative byte count to binary KiB."""

    number = _real("byte_count", byte_count)
    if number < 0.0:
        raise ValueError("byte_count must be non-negative")
    return number / 2**10


def bytes_to_mib(byte_count: Real) -> float:
    """Convert a finite non-negative byte count to binary MiB."""

    number = _real("byte_count", byte_count)
    if number < 0.0:
        raise ValueError("byte_count must be non-negative")
    return number / 2**20


def bytes_to_gib(byte_count: Real) -> float:
    """Convert a finite non-negative byte count to binary GiB."""

    number = _real("byte_count", byte_count)
    if number < 0.0:
        raise ValueError("byte_count must be non-negative")
    return number / 2**30



@dataclass(frozen=True)
class AffineTransform:
    """Elementwise affine state transform ``h_out = multiplier*h_in + bias``."""

    multiplier: np.ndarray
    bias: np.ndarray


def compose_affine(later: AffineTransform, earlier: AffineTransform) -> AffineTransform:
    """Compose two diagonal affine transforms as ``later(earlier(h))``."""

    a2 = _finite_array("later.multiplier", later.multiplier, ndim=1)
    b2 = _finite_array("later.bias", later.bias, ndim=1)
    a1 = _finite_array("earlier.multiplier", earlier.multiplier, ndim=1)
    b1 = _finite_array("earlier.bias", earlier.bias, ndim=1)
    if a1.shape != b1.shape or a2.shape != b2.shape or a1.shape != a2.shape:
        raise ValueError("all affine components must share shape [d_state]")
    return AffineTransform(multiplier=a2 * a1, bias=a2 * b1 + b2)


def affine_scan_sequential(
    multipliers: Any,
    biases: Any,
    *,
    initial_state: Any | None = None,
) -> np.ndarray:
    """Evaluate ``h_t = a_t*h_{t-1} + b_t`` sequentially.

    ``multipliers`` and ``biases`` must both have shape ``[T, d_state]``. The returned
    array contains post-update states with the same shape.
    """

    a = _finite_array("multipliers", multipliers, ndim=2)
    b = _finite_array("biases", biases, ndim=2)
    if a.shape != b.shape:
        raise ValueError("multipliers and biases must have identical shape")
    state = (
        np.zeros(a.shape[1], dtype=np.float64)
        if initial_state is None
        else _finite_array("initial_state", initial_state, ndim=1).copy()
    )
    if state.shape != (a.shape[1],):
        raise ValueError("initial_state must have shape [d_state]")
    result = np.empty_like(b)
    for index in range(a.shape[0]):
        state = a[index] * state + b[index]
        result[index] = state
    return result


def affine_scan_chunked(
    multipliers: Any,
    biases: Any,
    *,
    chunk_size: Integral,
    initial_state: Any | None = None,
) -> np.ndarray:
    """Evaluate the same affine recurrence in non-overlapping chunks.

    Inside each chunk, the local zero-state recurrence and cumulative multipliers can be
    computed in parallel by a scan kernel. Only one boundary state crosses each chunk.
    This NumPy version keeps the mechanics transparent; it is not a fused GPU scan.
    """

    a = _finite_array("multipliers", multipliers, ndim=2)
    b = _finite_array("biases", biases, ndim=2)
    if a.shape != b.shape:
        raise ValueError("multipliers and biases must have identical shape")
    chunk = _positive_int("chunk_size", chunk_size)
    state = (
        np.zeros(a.shape[1], dtype=np.float64)
        if initial_state is None
        else _finite_array("initial_state", initial_state, ndim=1).copy()
    )
    if state.shape != (a.shape[1],):
        raise ValueError("initial_state must have shape [d_state]")

    result = np.empty_like(b)
    for start in range(0, a.shape[0], chunk):
        end = min(start + chunk, a.shape[0])
        local_a = a[start:end]
        local_b = b[start:end]

        # Zero-start local states carry all within-chunk contributions.
        zero_state = np.zeros(a.shape[1], dtype=np.float64)
        local_zero = np.empty_like(local_b)
        for offset in range(local_a.shape[0]):
            zero_state = local_a[offset] * zero_state + local_b[offset]
            local_zero[offset] = zero_state

        # The chunk-start state is multiplied by the cumulative transition to each row.
        prefix_multiplier = np.cumprod(local_a, axis=0)
        result[start:end] = local_zero + prefix_multiplier * state
        state = result[end - 1].copy()
    return result


rng = np.random.default_rng(SEED)
a = rng.uniform(.3,.99,(53,6)); b = rng.normal(size=(53,6)); h0=rng.normal(size=6)
expected=affine_scan_sequential(a,b,initial_state=h0)
sizes=[1,7,16,64]
errors=[]
for chunk in sizes:
    actual=affine_scan_chunked(a,b,chunk_size=chunk,initial_state=h0)
    errors.append(np.max(np.abs(actual-expected)))
    np.testing.assert_allclose(actual,expected,rtol=1e-12,atol=1e-12)
# A state reset at each chunk is an intentionally wrong control.
reset=np.concatenate([affine_scan_sequential(a[j:j+16],b[j:j+16]) for j in range(0,53,16)])
assert np.max(np.abs(reset-expected)) > .1
print("chunks:",sizes,"max errors:",errors,"reset error:",np.max(np.abs(reset-expected)))
fig,ax=plt.subplots(figsize=(11,4),layout="constrained")
ax.plot(expected[:,0],color="#5E81AC",label="Эталон" if DEMO_LANGUAGE=="ru" else "Reference")
ax.plot(reset[:,0],"--",color="#D08770",label="Контроль со сбросом" if DEMO_LANGUAGE=="ru" else "Reset control")
for x in (16,32,48):ax.axvline(x,color="#D8DEE9")
ax.set_xlabel("Шаг" if DEMO_LANGUAGE=="ru" else "Step")
ax.set_ylabel("h[0]");ax.legend();plt.show()

"""Minimal M8 operators; used by the validated reference and embedded in notebooks."""
import numpy as np

def zoh_diagonal(a, b, delta):
    """Diagonal A: [n], B: [n,p], positive scalar step; caller validates shapes."""
    a = np.asarray(a, dtype=np.float64)
    b = np.asarray(b, dtype=np.float64)
    transition = np.exp(delta * a)
    factor = np.full_like(a, float(delta))
    np.divide(np.expm1(delta * a), a, out=factor, where=a != 0)
    return transition, factor[:, None] * b

def delta_step(state, key, value, beta, decay=1.0):
    """State [dv,dk], unit key [dk], value [dv]; decay before correction."""
    base = decay * state
    prediction = base @ key
    return base + beta * np.outer(value - prediction, key)

def compose_affine_pair(later, earlier):
    """Elementwise transforms; returns later(earlier(h)), not the reverse."""
    a2, b2 = later
    a1, b1 = earlier
    return a2 * a1, a2 * b1 + b2



@dataclass(frozen=True)
class DeltaUpdateResult:
    """State and retrieval trace for one normalized-key delta-rule update."""

    state: np.ndarray
    normalized_key: np.ndarray
    retrieval_before_write: np.ndarray
    retrieval_after_write: np.ndarray


def normalize_key(key: Any) -> np.ndarray:
    """Return a unit-L2 key, rejecting the zero vector."""

    vector = _finite_array("key", key, ndim=1)
    scale = float(np.max(np.abs(vector)))
    if scale == 0.0:
        raise ValueError("key must have non-zero norm")
    scaled = vector / scale
    return scaled / np.linalg.norm(scaled)


def delta_update(
    state: Any,
    key: Any,
    value: Any,
    *,
    beta: Real,
    decay: Real = 1.0,
    normalize: bool = True,
) -> DeltaUpdateResult:
    """Apply a gated delta-rule write to associative state ``[d_value, d_key]``.

    With a unit key, the update is

    ``base = decay * state``
    ``new = base + beta * outer(value - base@key, key)``.

    Thus ``beta=1`` overwrites retrieval at that key exactly, while ``beta=0`` only
    applies the global decay. If ``normalize=False``, the caller must supply a unit key;
    the explicit check prevents a misleading "exact overwrite" claim for arbitrary norm.
    """

    matrix = _finite_array("state", state, ndim=2)
    raw_key = _finite_array("key", key, ndim=1)
    target = _finite_array("value", value, ndim=1)
    if matrix.shape != (target.size, raw_key.size):
        raise ValueError("state must have shape [d_value, d_key]")
    write_rate = _unit_interval("beta", beta)
    forget_rate = _unit_interval("decay", decay)
    unit_key = normalize_key(raw_key) if normalize else raw_key.copy()
    if not normalize and not np.isclose(np.linalg.norm(unit_key), 1.0, rtol=0.0, atol=1e-10):
        raise ValueError("key must have unit norm when normalize=False")

    base = forget_rate * matrix
    before = base @ unit_key
    updated = delta_step(matrix, unit_key, target, write_rate, forget_rate)
    after = updated @ unit_key
    return DeltaUpdateResult(updated, unit_key, before, after)


# Repeated writes to the same four unit keys; compare on identical inputs.
rng=np.random.default_rng(SEED+1)
keys=rng.normal(size=(4,8));keys/=np.linalg.norm(keys,axis=1,keepdims=True)
values=rng.normal(size=(48,3));ids=np.arange(48)%4
states={name:np.zeros((3,8)) for name in ("additive","delta","gated")}
latest=np.zeros((4,3));history={name:[] for name in states}
for t,(i,v) in enumerate(zip(ids,values)):
    k=keys[i];latest[i]=v
    states["additive"] += np.outer(v,k)
    states["delta"]=delta_update(states["delta"],k,v,beta=1).state
    states["gated"]=delta_update(states["gated"],k,v,beta=1,decay=.9).state
    np.testing.assert_allclose(states["delta"]@k,v,atol=1e-12)
    np.testing.assert_allclose(states["gated"]@k,v,atol=1e-12)
    for name,state in states.items():
        history[name].append(np.mean((keys@state.T-latest)**2))
fig,ax=plt.subplots(figsize=(11,4.8),layout="constrained")
labels = {"additive":"Аддитивная запись", "delta":"Дельта-правило", "gated":"Дельта-правило с затуханием"}
for (name,errors),color in zip(history.items(),("#5E81AC","#A3BE8C","#D08770")):
    ax.plot(errors,label=labels[name] if DEMO_LANGUAGE=="ru" else name,color=color)
ax.set_xlabel("Запись" if DEMO_LANGUAGE=="ru" else "Write")
ax.set_ylabel("MSE относительно последних значений" if DEMO_LANGUAGE=="ru" else "MSE relative to latest values")
ax.legend();plt.show()
print({name:float(errors[-1]) for name,errors in history.items()})
print("Exact at the current key does not imply exact at all previous keys.")