Модуль 4. Позиционные представления: от синусоид к нескольким осям

import importlib.util
import subprocess
import sys

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

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

import numpy as np
import matplotlib.pyplot as plt

np.set_printoptions(precision=6, suppress=True)
rng = np.random.default_rng(42)

REFERENCE_JSON = '{"alibi":{"allowed_bias_with_null_for_future":[[[-0.0,null,null,null],[-1.0,-0.0,null,null],[-2.0,-1.0,-0.0,null],[-3.0,-2.0,-1.0,-0.0]],[[-0.0,null,null,null],[-0.5,-0.0,null,null],[-1.0,-0.5,-0.0,null],[-1.5,-1.0,-0.5,-0.0]]],"positions":[0,1,2,3],"slopes":[1.0,0.5]},"meta":{"description":"Frozen numerical contract independent of lecture prose and figures","module":"module-04","version":2},"mrope":{"axis_sections":[2,1,1],"base":10000.0,"head_dim":8,"image_coordinates":[[0,0,0],[0,0,1],[0,1,0],[0,1,1]],"image_input":[[0.0,0.1,0.2,0.3,0.4,0.5,0.6,0.7],[0.8,0.9,1.0,1.1,1.2,1.3,1.4,1.5],[1.6,1.7,1.8,1.9,2.0,2.1,2.2,2.3],[2.4,2.5,2.6,2.7,2.8,2.9,3.0,3.1]],"image_output":[[0.0,0.1,0.2,0.3,0.4,0.5,0.6,0.7],[0.8,0.9,1.0,1.1,1.2,1.3,1.3984993002500583,1.5013992497667292],[1.6,1.7,1.8,1.9,1.9789003508315806,2.1198946675433303,2.2,2.3],[2.4,2.5,2.6,2.7,2.7708604844975793,2.9278545345439957,2.996898500516792,3.1029984495001295]],"text_coordinates":[[4,4,4],[5,5,5],[6,6,6]],"text_input":[[0.0,0.1,0.2,0.3,0.4,0.5,0.6,0.7],[0.8,0.9,1.0,1.1,1.2,1.3,1.4,1.5],[1.6,1.7,1.8,1.9,2.0,2.1,2.2,2.3]],"text_matches_standard_1d":true,"text_output":[[0.07568024953079283,-0.06536436208636119,0.0673866961079819,0.3542019666625956,0.37968537557107407,0.5155957870051426,0.5971952074730607,0.7023943936074717],[1.0899615955674056,-0.5118434528136071,0.3502144694257494,1.444766356683613,1.1335273924220777,1.3583503416382703,1.3924825312864193,1.5069812208724322],[2.0112788055787596,1.185224690187341,0.4127834073868535,2.5844941204394525,1.8704766662635748,2.216149146822818,2.1861604829186514,2.3131585209243424]],"text_positions":[4,5,6],"text_standard_rope_output":[[0.07568024953079283,-0.06536436208636119,0.0673866961079819,0.3542019666625956,0.37968537557107407,0.5155957870051426,0.5971952074730607,0.7023943936074717],[1.0899615955674056,-0.5118434528136071,0.3502144694257494,1.444766356683613,1.1335273924220777,1.3583503416382703,1.3924825312864193,1.5069812208724322],[2.0112788055787596,1.185224690187341,0.4127834073868535,2.5844941204394525,1.8704766662635748,2.216149146822818,2.1861604829186514,2.3131585209243424]]},"scaling":{"base":150000.0,"beta_fast":32.0,"beta_slow":1.0,"extended_context":131072,"factor":32.0,"head_dim":64,"ntk":[1.0,0.6161605983892874,0.37965388300744474,0.23392776373468369,0.14413707088263056,0.08881158384512078,0.05472219864590999,0.033717662662841354,0.020775495202624464,0.01280104155588286,0.007887497425078919,0.00485996513323059,0.0029945190246424332,0.001845104634111787,0.0011368807754451658,0.0007005011388955706,0.00043162120091427207,0.00026594797743284077,0.00016386666491543985,0.00010096818231035427,6.221261563062656e-05,3.83329624743296e-05,2.3619261096217023e-05,1.4553258050557898e-05,8.967144188945469e-06,5.525200929303662e-06,3.4044111108207913e-06,2.0976639872064776e-06,1.2924978975768017e-06,7.963862779878181e-07,4.907018455939914e-07,3.0235114281192144e-07],"ntk_base":5367776.967415848,"original_context":4096,"position_interpolation":[0.03125,0.0215326345590051,0.014836939233653148,0.01022330849575739,0.007044312506339714,0.0048538434213901555,0.0033445131711816074,0.0023045177565706506,0.001587914837983171,0.0010941436773476293,0.0007539134706999197,0.0005194797841181623,0.00035794458727063204,0.00024663967968231775,0.00016994566689118156,0.0001171000940817352,8.068715304598785e-05,5.55970233646647e-05,3.8308812373753384e-05,2.639646903147278e-05,1.818833668168956e-05,1.2532569824095001e-05,8.63549587543848e-06,5.950239261491605e-06,4.09997848180298e-06,2.8250668271503398e-06,1.9465962110014815e-06,1.3412910350540445e-06,9.242089502428302e-07,6.36820914615699e-07,4.387978250864396e-07,3.0235114281192144e-07],"raw":[1.0,0.6890443058881632,0.4747820554769007,0.3271458718642365,0.22541800020287084,0.15532298948448497,0.10702442147781144,0.07374456821026082,0.050813274815461475,0.03501259767512414,0.02412523106239743,0.016623353091781194,0.011454226792660225,0.007892469749834168,0.00543826134051781,0.0037472030106155264,0.0025819888974716113,0.0017791047476692703,0.0012258819959601083,0.0008446870090071289,0.0005820267738140659,0.00040104223437104005,0.00027633586801403137,0.00019040765636773135,0.00013119931141769537,9.040213846881087e-05,6.229107875204741e-05,4.292131312172942e-05,2.9574686407770566e-05,2.0378269267702367e-05,1.4041530402766067e-05,9.675236569981486e-06],"yarn":{"attention_factor":1.3465735902799727,"hf_truncated_high":18,"hf_truncated_low":8,"high":17.39802450158856,"interpolated_pairs":14,"inverse_frequencies":[1.0,0.6890443058881632,0.4747820554769007,0.3271458718642365,0.22541800020287084,0.15532298948448497,0.10702442147781144,0.07374456821026082,0.050813274815461475,0.03170569618466377,0.01933500112654036,0.011592049256286924,0.006794959489732219,0.0038603593171920685,0.0020937923789696887,0.0010526021013863357,0.00045648391922324086,0.0001293187012450632,3.8308812373753384e-05,2.639646903147278e-05,1.818833668168956e-05,1.2532569824095001e-05,8.63549587543848e-06,5.950239261491605e-06,4.09997848180298e-06,2.8250668271503398e-06,1.9465962110014815e-06,1.3412910350540445e-06,9.242089502428302e-07,6.36820914615699e-07,4.387978250864396e-07,3.0235114281192144e-07],"kept_pairs":9,"low":8.092779115512402,"ramp":[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.09749564324709936,0.20496191184183662,0.3124281804365739,0.4198944490313112,0.5273607176260484,0.6348269862207857,0.742293254815523,0.8497595234102603,0.9572257920049976,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0],"score_factor_if_applied_to_q_and_k":1.8132604340394958,"transition_pairs":9,"truncate":false}},"sinusoidal":{"base":10000.0,"encoding":[[0.0,1.0,0.0,1.0,0.0,1.0,0.0,1.0],[0.8414709848078965,0.5403023058681398,0.09983341664682815,0.9950041652780258,0.009999833334166664,0.9999500004166653,0.0009999998333333417,0.9999995000000417],[0.9092974268256817,-0.4161468365471424,0.19866933079506122,0.9800665778412416,0.01999866669333308,0.9998000066665778,0.0019999986666669333,0.9999980000006666],[0.1411200080598672,-0.9899924966004454,0.2955202066613396,0.955336489125606,0.02999550020249566,0.9995500337489875,0.002999995500002025,0.999995500003375]],"model_dim":8,"positions":[0,1,2,3]},"spectrum":{"bases":{"10000":{"full_turn_counts":{"131072":32,"4096":23},"inverse_frequencies":[1.0,0.7498942093324558,0.5623413251903491,0.4216965034285822,0.31622776601683794,0.23713737056616552,0.1778279410038923,0.1333521432163324,0.1,0.07498942093324558,0.056234132519034905,0.042169650342858224,0.03162277660168379,0.023713737056616554,0.01778279410038923,0.01333521432163324,0.01,0.007498942093324558,0.005623413251903491,0.004216965034285822,0.0031622776601683794,0.0023713737056616554,0.0017782794100389228,0.001333521432163324,0.001,0.0007498942093324559,0.0005623413251903491,0.00042169650342858224,0.00031622776601683794,0.0002371373705661655,0.00017782794100389227,0.0001333521432163324],"maximum_wavelength_tokens":47117.2427801674,"minimum_wavelength_tokens":6.283185307179586,"wavelengths":[6.283185307179586,8.378762269377678,11.173259061216543,14.899780425245321,19.869176531592203,26.49597274431474,35.33294752055899,47.117242780167395,62.83185307179586,83.78762269377677,111.73259061216544,148.99780425245322,198.69176531592203,264.95972744314736,353.32947520558986,471.1724278016739,628.3185307179587,837.8762269377677,1117.3259061216543,1489.9780425245322,1986.9176531592202,2649.597274431474,3533.294752055899,4711.724278016739,6283.185307179586,8378.762269377676,11173.259061216542,14899.78042524532,19869.176531592202,26495.972744314742,35332.94752055899,47117.2427801674]},"150000":{"full_turn_counts":{"131072":27,"4096":18},"inverse_frequencies":[1.0,0.6890443058881632,0.4747820554769007,0.3271458718642365,0.22541800020287084,0.15532298948448497,0.10702442147781144,0.07374456821026082,0.050813274815461475,0.03501259767512414,0.02412523106239743,0.016623353091781194,0.011454226792660225,0.007892469749834168,0.00543826134051781,0.0037472030106155264,0.0025819888974716113,0.0017791047476692703,0.0012258819959601083,0.0008446870090071289,0.0005820267738140659,0.00040104223437104005,0.00027633586801403137,0.00019040765636773135,0.00013119931141769537,9.040213846881087e-05,6.229107875204741e-05,4.292131312172942e-05,2.9574686407770566e-05,2.0378269267702367e-05,1.4041530402766067e-05,9.675236569981486e-06],"maximum_wavelength_tokens":649408.9588128396,"minimum_wavelength_tokens":6.283185307179586,"wavelengths":[6.283185307179586,9.118695638999144,13.233830627546281,19.206066307286523,27.873485265262175,40.452384595695705,58.70795861748462,85.20200822472704,123.65243787176136,179.4549883296344,260.4404198628719,377.9734011837886,548.5473110420513,796.0987506238595,1155.3665618028806,1676.7667215733509,2433.467205584167,3531.6556348977883,5125.4405627016395,7438.477495427609,10795.354423311823,15667.141185350692,22737.494601535214,32998.59589177952,47890.3832595279,69502.61811945203,100868.14088081702,146388.46880941882,212451.46002727243,308327.72031027393,447471.54526274785,649408.9588128396]},"500000":{"full_turn_counts":{"131072":25,"4096":16},"inverse_frequencies":[1.0,0.6636012376960885,0.44036660267178046,0.2922278225730151,0.19392274474868576,0.12868737343265052,0.08539710028576561,0.05666962144529105,0.03760603093086393,0.024955408670558694,0.016560440080994446,0.010989528534539826,0.007292664737217109,0.004839421345719893,0.003211445994752591,0.0021311195369119653,0.001414213562373095,0.0009384738703573802,0.000622772421914596,0.0004132725499855165,0.0002742481756762073,0.00018199142881462546,0.00012076973741146504,8.01429472224798e-05,5.318295896944988e-05,3.529227739646723e-05,2.341999896140934e-05,1.5541540297632344e-05,1.031338537721246e-05,6.8439753011549275e-06,4.5416704806078695e-06,3.013858152139171e-06],"maximum_wavelength_tokens":2084764.7732591254,"minimum_wavelength_tokens":6.283185307179586,"wavelengths":[6.283185307179586,9.468314629722128,14.268078616903308,21.500982527458316,32.40045573469106,48.8251888245115,73.57609668424418,110.87395939719457,167.07919319459117,251.77649423117703,379.40931982782695,571.7429357803371,861.5760539647758,1298.3340069648193,1956.497265545218,2948.3026167007256,4442.882938158366,6695.109481084312,10089.055144515105,15203.490547339246,22910.58196353461,34524.62211052572,52026.15689866611,78399.72854675821,118142.83050307268,178032.86641424094,268282.9028956321,404283.3069857812,609226.2702664398,918060.7805699259,1383452.483839961,2084764.7732591254]}},"contexts":[4096,131072],"head_dim":64},"worked_rope":{"base":10000.0,"head_dim":4,"k":[0.0,1.0,1.0,-0.5],"logits":[1.1442484284723262,1.1442484284723262,1.1442484284723264],"placements":[[3,1],[7,5],[2,0]],"q":[1.0,0.0,0.5,0.5],"rotated_k_at_1":[-0.8414709848078965,0.5403023058681398,1.0049499170837486,-0.489975166874166],"rotated_q_at_3":[-0.9899924966004454,0.1411200080598672,0.48477726677324595,0.5147727669757416]}}'
REFERENCE_SHA256 = '88813060dd14b01e5a13f3dfe2736dfd6b07c38649bca8e0b28200d6bca8ceb3'
reference = json.loads(REFERENCE_JSON)
actual_sha256 = hashlib.sha256(
    json.dumps(reference, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
).hexdigest()
assert actual_sha256 == REFERENCE_SHA256
print('Встроенный контракт М4 проверен', actual_sha256)

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

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


def _positive_even_int(name: str, value: object, *, minimum: int = 2) -> int:
    """Return an even integer not smaller than ``minimum``."""

    result = _positive_int(name, value)
    if result < minimum:
        raise ValueError(f"{name} must be at least {minimum}")
    if result % 2:
        raise ValueError(f"{name} must be even")
    return result


def _finite_real(name: str, value: object) -> float:
    """Return a finite real scalar and reject booleans and complex values."""

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


def _real_greater_than(name: str, value: object, lower: float) -> float:
    """Return a finite scalar strictly greater than ``lower``."""

    result = _finite_real(name, value)
    if result <= lower:
        raise ValueError(f"{name} must be greater than {lower}")
    return result


def _finite_array(
    name: str,
    value: object,
    *,
    ndim: int | None = None,
    minimum_ndim: int | None = None,
) -> np.ndarray:
    """Convert to float64 and enforce rank, non-emptiness, and finite values."""

    try:
        if np.iscomplexobj(value):
            raise TypeError(f"{name} must be real")
        array = np.asarray(value, dtype=np.float64)
    except (TypeError, ValueError) as exc:
        raise TypeError(f"{name} must be numeric") from exc
    if ndim is not None and array.ndim != ndim:
        raise ValueError(f"{name} must have rank {ndim}, got shape {array.shape}")
    if minimum_ndim is not None and array.ndim < minimum_ndim:
        raise ValueError(f"{name} must have rank at least {minimum_ndim}, got shape {array.shape}")
    if 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 _position_vector(name: str, value: object) -> np.ndarray:
    """Validate a non-empty one-dimensional vector of non-negative integer positions."""

    raw = np.asarray(value)
    if raw.ndim != 1 or raw.size == 0:
        raise ValueError(f"{name} must be a non-empty one-dimensional vector")
    if raw.dtype.kind not in "iu":
        raise TypeError(f"{name} must contain integer positions")
    result = raw.astype(np.int64, copy=False)
    if np.any(result < 0):
        raise ValueError(f"{name} must contain non-negative positions")
    return result


def _inverse_frequency_vector(value: object, *, expected_pairs: int | None = None) -> np.ndarray:
    """Validate a positive one-dimensional inverse-frequency vector."""

    inv = _finite_array("inverse_frequencies", value, ndim=1)
    if np.any(inv <= 0.0):
        raise ValueError("inverse_frequencies must be strictly positive")
    if expected_pairs is not None and inv.size != expected_pairs:
        raise ValueError(
            f"inverse_frequencies must have {expected_pairs} entries, got {inv.size}"
        )
    return inv


def _sequence_tensor(name: str, value: object) -> np.ndarray:
    """Validate a finite sequence tensor with shape ``[..., sequence, head_dim]``."""

    tensor = _finite_array(name, value, minimum_ndim=2)
    _positive_even_int(f"{name}.shape[-1]", tensor.shape[-1])
    if tensor.shape[-2] <= 0:
        raise ValueError(f"{name} must have a non-empty sequence axis")
    return tensor

def rope_inverse_frequencies(head_dim: int, base: Real = 10_000.0) -> np.ndarray:
    """Return ``base**(-2i/head_dim)`` for the ``head_dim/2`` RoPE pairs."""

    dim = _positive_even_int("head_dim", head_dim)
    rope_base = _real_greater_than("base", base, 1.0)
    pair_indices = np.arange(dim // 2, dtype=np.float64)
    inv = rope_base ** (-2.0 * pair_indices / dim)
    if not np.all(np.isfinite(inv)) or np.any(inv <= 0.0):
        raise FloatingPointError("inverse-frequency construction produced invalid values")
    return inv


def sinusoidal_encoding(
    positions: object,
    model_dim: int,
    *,
    base: Real = 10_000.0,
) -> np.ndarray:
    """Return the classic sinusoidal encoding with shape ``[positions, model_dim]``."""

    pos = _position_vector("positions", positions)
    dim = _positive_even_int("model_dim", model_dim)
    inv = rope_inverse_frequencies(dim, base)
    angles = pos.astype(np.float64)[:, None] * inv[None, :]
    result = np.empty((pos.size, dim), dtype=np.float64)
    result[:, 0::2] = np.sin(angles)
    result[:, 1::2] = np.cos(angles)
    if not np.all(np.isfinite(result)):
        raise FloatingPointError("sinusoidal encoding produced non-finite values")
    return result


def wavelengths_from_inverse_frequencies(inverse_frequencies: object) -> np.ndarray:
    """Convert radians-per-position frequencies to tokens per full ``2π`` turn."""

    inv = _inverse_frequency_vector(inverse_frequencies)
    return (2.0 * math.pi) / inv


def rotations_within_context(
    inverse_frequencies: object,
    context_length: int,
) -> np.ndarray:
    """Return turns (including fractions) across a span L = context_length.

    This is the YaRN L*theta/(2*pi) diagnostic. Positions 0..L-1 have a maximum
    endpoint separation of L-1, which is a different convention.
    """

    inv = _inverse_frequency_vector(inverse_frequencies)
    length = _positive_int("context_length", context_length)
    return length * inv / (2.0 * math.pi)


def full_turn_count(inverse_frequencies: object, context_length: int) -> int:
    """Count frequency pairs that complete at least one full turn in the context."""

    return int(np.count_nonzero(rotations_within_context(inverse_frequencies, context_length) >= 1.0))


positions = np.array(reference["sinusoidal"]["positions"], dtype=np.int64)
table = sinusoidal_encoding(
    positions,
    reference["sinusoidal"]["model_dim"],
    base=reference["sinusoidal"]["base"],
)
assert table.shape == (4, 8)
assert np.allclose(table, np.array(reference["sinusoidal"]["encoding"]), atol=1e-12)
assert np.allclose(table[:, 0] ** 2 + table[:, 1] ** 2, 1.0, atol=1e-12)

for base, expected in reference["spectrum"]["bases"].items():
    inv = rope_inverse_frequencies(reference["spectrum"]["head_dim"], float(base))
    wave = wavelengths_from_inverse_frequencies(inv)
    assert np.allclose(wave, np.array(expected["wavelengths"]), atol=1e-9)
    for context, count in expected["full_turn_counts"].items():
        assert full_turn_count(inv, int(context)) == count
print('Синусоидальная таблица и спектр проверены')

def rotate_half_split(x: object) -> np.ndarray:
    """Apply the split-half quarter turn used by HF/Llama-style RoPE implementations."""

    values = _sequence_tensor("x", x)
    half = values.shape[-1] // 2
    return np.concatenate((-values[..., half:], values[..., :half]), axis=-1)


def _angles_for_sequence(
    x: np.ndarray,
    positions: np.ndarray,
    inverse_frequencies: np.ndarray,
) -> np.ndarray:
    """Build a broadcastable angle tensor for ``[..., T, D]`` sequence tensors."""

    if positions.size != x.shape[-2]:
        raise ValueError(
            f"positions length {positions.size} does not match sequence axis {x.shape[-2]}"
        )
    if inverse_frequencies.size != x.shape[-1] // 2:
        raise ValueError("inverse_frequencies do not match head_dim")
    base_angles = positions.astype(np.float64)[:, None] * inverse_frequencies[None, :]
    return base_angles.reshape((1,) * (x.ndim - 2) + base_angles.shape)


def apply_rope_adjacent(
    x: object,
    positions: object,
    inverse_frequencies: object,
) -> np.ndarray:
    """Apply RoPE to adjacent pairs ``(x_0,x_1), (x_2,x_3), ...``."""

    values = _sequence_tensor("x", x)
    pos = _position_vector("positions", positions)
    inv = _inverse_frequency_vector(
        inverse_frequencies,
        expected_pairs=values.shape[-1] // 2,
    )
    angles = _angles_for_sequence(values, pos, inv)
    cosine = np.cos(angles)
    sine = np.sin(angles)
    first = values[..., 0::2]
    second = values[..., 1::2]
    result = np.empty_like(values, dtype=np.float64)
    result[..., 0::2] = cosine * first - sine * second
    result[..., 1::2] = sine * first + cosine * second
    if not np.all(np.isfinite(result)):
        raise FloatingPointError("adjacent-pair RoPE produced non-finite values")
    return result


def apply_rope_split_half(
    x: object,
    positions: object,
    inverse_frequencies: object,
) -> np.ndarray:
    """Apply split-half RoPE to pairs ``(x_j, x_{j+D/2})``."""

    values = _sequence_tensor("x", x)
    pos = _position_vector("positions", positions)
    inv = _inverse_frequency_vector(
        inverse_frequencies,
        expected_pairs=values.shape[-1] // 2,
    )
    angles = _angles_for_sequence(values, pos, inv)
    repeated = np.concatenate((angles, angles), axis=-1)
    result = values * np.cos(repeated) + rotate_half_split(values) * np.sin(repeated)
    if not np.all(np.isfinite(result)):
        raise FloatingPointError("split-half RoPE produced non-finite values")
    return result


def apply_rope_complex_adjacent(
    x: object,
    positions: object,
    inverse_frequencies: object,
) -> np.ndarray:
    """Apply adjacent-pair RoPE through the equivalent complex multiplication."""

    values = _sequence_tensor("x", x)
    pos = _position_vector("positions", positions)
    inv = _inverse_frequency_vector(
        inverse_frequencies,
        expected_pairs=values.shape[-1] // 2,
    )
    angles = _angles_for_sequence(values, pos, inv)
    complex_values = values[..., 0::2] + 1j * values[..., 1::2]
    rotated = complex_values * np.exp(1j * angles)
    result = np.empty_like(values, dtype=np.float64)
    result[..., 0::2] = rotated.real
    result[..., 1::2] = rotated.imag
    if not np.all(np.isfinite(result)):
        raise FloatingPointError("complex RoPE produced non-finite values")
    return result


def adjacent_to_split_permutation(head_dim: int) -> np.ndarray:
    """Return the coordinate permutation mapping adjacent pairs to split-half pairs."""

    dim = _positive_even_int("head_dim", head_dim)
    return np.concatenate((np.arange(0, dim, 2), np.arange(1, dim, 2)))


def rope_attention_logits(
    queries: object,
    keys: object,
    query_positions: object,
    key_positions: object,
    inverse_frequencies: object,
    *,
    pairing: str = "split_half",
) -> np.ndarray:
    """Return unscaled rotated dot products ``[..., query, key]``.

    Divide by sqrt(head_dim) before the standard attention softmax. The historical
    function name and fixture key 'logits' are retained for file compatibility.
    """

    q = _sequence_tensor("queries", queries)
    k = _sequence_tensor("keys", keys)
    if q.ndim != k.ndim or q.shape[:-2] != k.shape[:-2]:
        raise ValueError("queries and keys must have identical batch dimensions")
    if q.shape[-1] != k.shape[-1]:
        raise ValueError("queries and keys must have the same head_dim")
    if pairing == "split_half":
        apply = apply_rope_split_half
    elif pairing == "adjacent":
        apply = apply_rope_adjacent
    else:
        raise ValueError("pairing must be 'split_half' or 'adjacent'")
    q_rot = apply(q, query_positions, inverse_frequencies)
    k_rot = apply(k, key_positions, inverse_frequencies)
    return np.matmul(q_rot, np.swapaxes(k_rot, -1, -2))


spec_inv = rope_inverse_frequencies(64, 150_000.0)
positions = np.array([0, 3, 17, 128, 4095], dtype=np.int64)
x = rng.standard_normal((2, 4, positions.size, 64))

adjacent = apply_rope_adjacent(x, positions, spec_inv)
complex_form = apply_rope_complex_adjacent(x, positions, spec_inv)
assert np.allclose(adjacent, complex_form, atol=1e-12)
assert np.allclose(np.linalg.norm(adjacent, axis=-1), np.linalg.norm(x, axis=-1), atol=1e-12)

perm = adjacent_to_split_permutation(64)
split = apply_rope_split_half(x[..., perm], positions, spec_inv)
assert np.allclose(split, adjacent[..., perm], atol=1e-12)

q = rng.standard_normal((1, 1, 64))
k = rng.standard_normal((1, 1, 64))
raw_adjacent = rope_attention_logits(q, k, [3], [1], spec_inv, pairing="adjacent").item()
raw_split = rope_attention_logits(q, k, [3], [1], spec_inv, pairing="split_half").item()
assert abs(raw_adjacent - raw_split) > 1e-3
q_perm, k_perm = q[..., perm], k[..., perm]
permuted_split = rope_attention_logits(q_perm, k_perm, [3], [1], spec_inv, pairing="split_half").item()
assert abs(raw_adjacent - permuted_split) < 1e-10
print('Три реализации, перестановка и нормы проверены', f"raw convention gap={abs(raw_adjacent-raw_split):.6f}")

max_error = 0.0
for _ in range(200):
    query = rng.standard_normal((1, 1, 64))
    key = rng.standard_normal((1, 1, 64))
    m, n, shift = (int(value) for value in rng.integers(0, 100_000, size=3))
    first = rope_attention_logits(query, key, [m], [n], spec_inv, pairing="split_half").item()
    second = rope_attention_logits(
        query, key, [m + shift], [n + shift], spec_inv, pairing="split_half"
    ).item()
    max_error = max(max_error, abs(first - second))
assert max_error < 1e-9

worked = reference["worked_rope"]
q4 = np.array(worked["q"], dtype=np.float64)[None, :]
k4 = np.array(worked["k"], dtype=np.float64)[None, :]
inv4 = rope_inverse_frequencies(worked["head_dim"], worked["base"])
observed = []
for q_pos, k_pos in worked["placements"]:
    observed.append(
        float(rope_attention_logits(q4, k4, [q_pos], [k_pos], inv4, pairing="adjacent").item())
    )
assert np.allclose(observed, worked["logits"], atol=1e-12)
assert np.float32(100_000_000) == np.float32(100_000_001)
large_positions = np.array([100_000_000, 100_000_001], dtype=np.int64)
large_x = np.ones((2, 4), dtype=np.float64)
large_rotated = apply_rope_adjacent(large_x, large_positions, inv4)
assert not np.allclose(large_rotated[0], large_rotated[1], atol=0.0, rtol=0.0)
print('Общий сдвиг и большие позиции проверены', f"max_error={max_error:.3e}")

def position_interpolated_frequencies(
    inverse_frequencies: object,
    factor: Real,
) -> np.ndarray:
    """Apply position interpolation by dividing every frequency by ``factor``."""

    inv = _inverse_frequency_vector(inverse_frequencies)
    scale = _finite_real("factor", factor)
    if scale < 1.0:
        raise ValueError("factor must be at least 1")
    return inv / scale


def ntk_scaled_base(head_dim: int, base: Real, factor: Real) -> float:
    """Return the static NTK-aware base ``base * factor**(d/(d-2))``."""

    dim = _positive_even_int("head_dim", head_dim, minimum=4)
    rope_base = _real_greater_than("base", base, 1.0)
    scale = _finite_real("factor", factor)
    if scale < 1.0:
        raise ValueError("factor must be at least 1")
    result = rope_base * scale ** (dim / (dim - 2))
    if not math.isfinite(result):
        raise FloatingPointError("NTK base is non-finite")
    return result


def ntk_scaled_frequencies(head_dim: int, base: Real, factor: Real) -> np.ndarray:
    """Return inverse frequencies generated from the static NTK-aware base."""

    return rope_inverse_frequencies(head_dim, ntk_scaled_base(head_dim, base, factor))


@dataclass(frozen=True)
class YarnSpectrum:
    """Piecewise YaRN spectrum plus diagnostic boundaries and scaling metadata."""

    inverse_frequencies: np.ndarray
    original_inverse_frequencies: np.ndarray
    ramp: np.ndarray
    low: float
    high: float
    truncate: bool
    kept_pairs: int
    transition_pairs: int
    interpolated_pairs: int
    attention_factor: float
    score_factor_if_applied_to_q_and_k: float


def yarn_scaled_frequencies(
    head_dim: int,
    base: Real,
    factor: Real,
    original_context: int,
    *,
    beta_fast: Real = 32.0,
    beta_slow: Real = 1.0,
    truncate: bool = False,
) -> YarnSpectrum:
    """Construct the lecture's piecewise YaRN frequency spectrum.

    ``ramp == 0`` keeps a frequency, ``ramp == 1`` divides it by ``factor``.
    With ``truncate=False`` this follows the continuous correction bounds used by the
    original gpt-oss implementation. ``truncate=True`` reproduces the integer-boundary
    option supported by the generic Hugging Face YaRN utility.
    """

    dim = _positive_even_int("head_dim", head_dim, minimum=4)
    rope_base = _real_greater_than("base", base, 1.0)
    scale = _finite_real("factor", factor)
    if scale < 1.0:
        raise ValueError("factor must be at least 1")
    context = _positive_int("original_context", original_context)
    fast = _real_greater_than("beta_fast", beta_fast, 0.0)
    slow = _real_greater_than("beta_slow", beta_slow, 0.0)
    if fast <= slow:
        raise ValueError("beta_fast must be greater than beta_slow")
    if not isinstance(truncate, bool):
        raise TypeError("truncate must be a bool")

    original = rope_inverse_frequencies(dim, rope_base)
    pair_count = dim // 2

    def correction_dimension(rotations: float) -> float:
        """Map a target turn count to the corresponding continuous pair index."""

        return dim * math.log(context / (rotations * 2.0 * math.pi)) / (2.0 * math.log(rope_base))

    low_raw = correction_dimension(fast)
    high_raw = correction_dimension(slow)
    low = math.floor(low_raw) if truncate else low_raw
    high = math.ceil(high_raw) if truncate else high_raw
    low = max(0.0, min(float(pair_count - 1), float(low)))
    high = max(0.0, min(float(pair_count - 1), float(high)))
    if high <= low:
        raise ValueError("YaRN correction range is empty for this configuration")
    pair_indices = np.arange(pair_count, dtype=np.float64)
    ramp = np.clip((pair_indices - low) / (high - low), 0.0, 1.0)
    scaled = original * (1.0 - ramp) + (original / scale) * ramp
    attention_factor = 0.1 * math.log(scale) + 1.0 if scale > 1.0 else 1.0
    score_factor_if_applied_to_q_and_k = attention_factor**2
    kept = int(np.count_nonzero(ramp == 0.0))
    interpolated = int(np.count_nonzero(ramp == 1.0))
    transition = pair_count - kept - interpolated
    return YarnSpectrum(
        inverse_frequencies=scaled,
        original_inverse_frequencies=original,
        ramp=ramp,
        low=low,
        high=high,
        truncate=truncate,
        kept_pairs=kept,
        transition_pairs=transition,
        interpolated_pairs=interpolated,
        attention_factor=attention_factor,
        score_factor_if_applied_to_q_and_k=score_factor_if_applied_to_q_and_k,
    )


scaling = reference["scaling"]
raw = np.array(scaling["raw"], dtype=np.float64)
pi = position_interpolated_frequencies(raw, scaling["factor"])
ntk = ntk_scaled_frequencies(
    scaling["head_dim"], scaling["base"], scaling["factor"]
)
yarn = yarn_scaled_frequencies(
    scaling["head_dim"],
    scaling["base"],
    scaling["factor"],
    scaling["original_context"],
    beta_fast=32.0,
    beta_slow=1.0,
)
assert np.allclose(pi, scaling["position_interpolation"], atol=1e-12)
assert np.allclose(ntk, scaling["ntk"], atol=1e-12)
assert math.isclose(
    ntk_scaled_base(scaling["head_dim"], scaling["base"], scaling["factor"]),
    scaling["ntk_base"],
    rel_tol=1e-12,
)
# PI can be read either as position/factor or inverse_frequency/factor.
probe_positions = np.array([0, 1, 17, 4096, 131071], dtype=np.int64)
angles_by_position = np.outer(probe_positions.astype(np.float64) / scaling["factor"], raw)
angles_by_frequency = np.outer(probe_positions.astype(np.float64), pi)
assert np.allclose(angles_by_position, angles_by_frequency, atol=1e-15)
assert np.allclose(yarn.inverse_frequencies, scaling["yarn"]["inverse_frequencies"], atol=1e-12)
assert abs(yarn.low - 8.092779115512402) < 1e-12
assert abs(yarn.high - 17.39802450158856) < 1e-12
assert yarn.truncate is False
assert (yarn.kept_pairs, yarn.transition_pairs, yarn.interpolated_pairs) == (9, 9, 14)
yarn_truncated = yarn_scaled_frequencies(
    scaling["head_dim"], scaling["base"], scaling["factor"], scaling["original_context"],
    beta_fast=scaling["beta_fast"], beta_slow=scaling["beta_slow"], truncate=True,
)
assert (yarn_truncated.low, yarn_truncated.high) == (8.0, 18.0)
assert abs(yarn.attention_factor - 1.34657359028) < 1e-12
assert abs(yarn.score_factor_if_applied_to_q_and_k - 1.81326043404) < 1e-12
# In this declared scenario, the chosen beta boundaries happen to make the fully
# interpolated pairs coincide with pairs below one full turn. This is a checked
# property of the scenario, not a general theorem about extrapolation.
no_turn = np.flatnonzero(rotations_within_context(raw, scaling["original_context"]) < 1.0)
fully_interpolated = np.flatnonzero(yarn.ramp == 1.0)
assert np.array_equal(no_turn, fully_interpolated)
assert np.allclose(position_interpolated_frequencies(raw, 1.0), raw)
print(
    'PI/NTK-aware/YaRN и масштабы проверены',
    f"bounds={yarn.low:.4f}/{yarn.high:.4f}; "
    f"zones={yarn.kept_pairs}/{yarn.transition_pairs}/{yarn.interpolated_pairs}",
)

from matplotlib.colors import LinearSegmentedColormap

INK = '#2E3440'
MUTED = '#4C566A'
BORDER = '#D8DEE9'
BLUE = '#5E81AC'
ORANGE = '#D08770'
PURPLE = '#B48EAD'
GREEN = '#A3BE8C'
PALE = {'blue': '#E8EEF5', 'green': '#EFF4EA', 'purple': '#F3EEF3', 'orange': '#FAF0E8', 'yellow': '#FAF5E3', 'gray': '#ECEFF4'}

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


def scaling_plot(reference, lang='en'):
    plot_style()
    ru = lang == 'ru'; sc = reference['scaling']; yr = sc['yarn']; x = np.arange(len(sc['raw']))
    fig, (ax, ramp_ax) = plt.subplots(2, 1, figsize=(11, 7), sharex=True,
                                    height_ratios=[3, 1], layout='constrained')
    series = [('Исходные' if ru else 'Original', sc['raw'], MUTED, '-'),
              ('PI', sc['position_interpolation'], PURPLE, '--'),
              ('Статическая NTK-aware' if ru else 'Static NTK-aware', sc['ntk'], ORANGE, '-.'),
              ('YaRN', yr['inverse_frequencies'], BLUE, '-')]
    for name, values, color, style in series:
        ax.plot(x, values, linestyle=style, color=color, linewidth=2.2, label=name)
    ax.set_yscale('log'); ax.grid(axis='y')
    ax.set_ylabel('Частота, рад/шаг (лог. шкала)' if ru else 'Frequency, rad/step (log scale)')
    ax.set_title('Как меняется спектр · s=32' if ru else 'How the spectrum changes · s=32', pad=16)
    ax.legend(loc='lower left', ncol=2)
    ramp_ax.plot(x, yr['ramp'], color=BLUE, marker='o', markersize=3)
    for value in (yr['low'], yr['high']):
        ramp_ax.axvline(value, color=MUTED, linestyle=':', linewidth=1.4)
    ramp_ax.set(xlabel='Индекс пары i' if ru else 'Pair index i',
                ylabel='Вес PI' if ru else 'PI weight', ylim=(-.08,1.08), yticks=[0,.5,1])
    ramp_ax.grid(axis='y')
    return fig


scaling_plot(reference, 'ru')
plt.show()

def alibi_bias(
    slopes: object,
    query_positions: object,
    key_positions: object,
    *,
    causal: bool = True,
) -> np.ndarray:
    """Return an ALiBi tensor with shape ``[heads, query, key]``.

    In causal mode, future keys receive ``-inf`` and allowed keys receive
    ``-slope * (query_position - key_position)``. In bidirectional mode the absolute
    distance is used.
    """

    if not isinstance(causal, bool):
        raise TypeError("causal must be a bool")
    slope_values = _finite_array("slopes", slopes, ndim=1)
    if np.any(slope_values <= 0.0):
        raise ValueError("slopes must be strictly positive")
    q_pos = _position_vector("query_positions", query_positions)
    k_pos = _position_vector("key_positions", key_positions)
    delta = q_pos[:, None] - k_pos[None, :]
    if causal:
        allowed = delta >= 0
        base = -np.abs(delta).astype(np.float64)
        bias = slope_values[:, None, None] * base[None, :, :]
        bias = np.where(allowed[None, :, :], bias, -np.inf)
    else:
        bias = -slope_values[:, None, None] * np.abs(delta)[None, :, :]
    return bias


def apply_mrope_adjacent(
    x: object,
    coordinates: object,
    axis_sections: Sequence[int],
    inverse_frequencies: object,
) -> np.ndarray:
    """Apply adjacent-pair multi-axis RoPE.

    ``coordinates`` has shape ``[sequence, axes]``. ``axis_sections`` gives the number of
    frequency pairs assigned to each axis and must sum to ``head_dim / 2``.
    """

    values = _sequence_tensor("x", x)
    raw_coordinates = np.asarray(coordinates)
    if raw_coordinates.ndim != 2 or raw_coordinates.shape[0] != values.shape[-2]:
        raise ValueError("coordinates must have shape [sequence, axes]")
    if raw_coordinates.dtype.kind not in "iu":
        raise TypeError("coordinates must contain integer positions")
    coords = raw_coordinates.astype(np.int64, copy=False)
    if np.any(coords < 0):
        raise ValueError("coordinates must be non-negative")

    if isinstance(axis_sections, (str, bytes)):
        raise TypeError("axis_sections must be a sequence of integers")
    sections: list[int] = []
    for index, value in enumerate(axis_sections):
        sections.append(_positive_int(f"axis_sections[{index}]", value))
    if len(sections) != coords.shape[1]:
        raise ValueError("axis_sections length must match the coordinate axis count")
    pair_count = values.shape[-1] // 2
    if sum(sections) != pair_count:
        raise ValueError("axis_sections must sum to head_dim / 2")

    inv = _inverse_frequency_vector(inverse_frequencies, expected_pairs=pair_count)
    axis_index = np.repeat(np.arange(len(sections)), sections)
    base_angles = coords[:, axis_index].astype(np.float64) * inv[None, :]
    angles = base_angles.reshape((1,) * (values.ndim - 2) + base_angles.shape)
    cosine = np.cos(angles)
    sine = np.sin(angles)
    first = values[..., 0::2]
    second = values[..., 1::2]
    result = np.empty_like(values, dtype=np.float64)
    result[..., 0::2] = cosine * first - sine * second
    result[..., 1::2] = sine * first + cosine * second
    if not np.all(np.isfinite(result)):
        raise FloatingPointError("M-RoPE produced non-finite values")
    return result


slopes = np.array(reference["alibi"]["slopes"], dtype=np.float64)
bias = alibi_bias(slopes, np.arange(4), np.arange(4), causal=True)
assert bias.shape == (2, 4, 4)
assert np.allclose(np.diagonal(bias, axis1=1, axis2=2), 0.0)
future_mask = np.triu(np.ones((4, 4), dtype=bool), k=1)
assert np.all(np.isneginf(bias[0][future_mask]))
assert np.all(np.isfinite(bias[0][~future_mask]))
rectangular_bias = alibi_bias(slopes, np.array([4, 5, 6]), np.arange(7), causal=True)
assert rectangular_bias.shape == (2, 3, 7)
assert np.all(np.isneginf(rectangular_bias[:, 0, 5:]))
assert np.all(np.isfinite(rectangular_bias[:, 2, :]))

mref = reference["mrope"]
m_inv = rope_inverse_frequencies(mref["head_dim"], mref["base"])
image_input = np.array(mref["image_input"], dtype=np.float64)
image_coordinates = np.array(mref["image_coordinates"], dtype=np.int64)
image_output = apply_mrope_adjacent(
    image_input, image_coordinates, mref["axis_sections"], m_inv
)
assert np.allclose(image_output, mref["image_output"], atol=1e-12)
assert np.all(image_coordinates[:, 0] == image_coordinates[0, 0])
assert np.allclose(
    np.linalg.norm(image_output, axis=-1), np.linalg.norm(image_input, axis=-1), atol=1e-12
)

text_input = np.array(mref["text_input"], dtype=np.float64)
text_positions = np.array(mref["text_positions"], dtype=np.int64)
text_coordinates = np.array(mref["text_coordinates"], dtype=np.int64)
text_output = apply_mrope_adjacent(
    text_input, text_coordinates, mref["axis_sections"], m_inv
)
text_1d = apply_rope_adjacent(text_input, text_positions, m_inv)
assert np.all(text_coordinates == text_positions[:, None])
assert np.allclose(text_output, mref["text_output"], atol=1e-12)
assert np.allclose(text_output, text_1d, atol=1e-12)

query = rng.standard_normal((1, 8))
key = rng.standard_normal((1, 8))
q_coord = np.array([[5, 1, 2]], dtype=np.int64)
k_coord = np.array([[8, 3, 1]], dtype=np.int64)
shift = np.array([[11, 4, 9]], dtype=np.int64)
first = float(
    apply_mrope_adjacent(query, q_coord, (2, 1, 1), m_inv)[0]
    @ apply_mrope_adjacent(key, k_coord, (2, 1, 1), m_inv)[0]
)
second = float(
    apply_mrope_adjacent(query, q_coord + shift, (2, 1, 1), m_inv)[0]
    @ apply_mrope_adjacent(key, k_coord + shift, (2, 1, 1), m_inv)[0]
)
assert abs(first - second) < 1e-12
print('Геометрия ALiBi и M-RoPE проверена')

def expect_error(expected_exception, function, *args, **kwargs):
    """Assert that one invalid call fails with the declared exception type."""

    try:
        function(*args, **kwargs)
    except expected_exception:
        return
    raise AssertionError(f"{function.__name__} did not raise {expected_exception.__name__}")

expect_error(ValueError, rope_inverse_frequencies, 63, 10_000.0)
expect_error(ValueError, rope_inverse_frequencies, 64, 1.0)
expect_error(ValueError, apply_rope_adjacent, np.ones((2, 8)), [0], np.ones(4))
expect_error(ValueError, apply_rope_adjacent, np.array([[np.nan] * 8]), [0], np.ones(4))
expect_error(ValueError, rope_attention_logits, np.ones((1, 1, 8)), np.ones((1, 1, 6)), [0], [0], np.ones(4))
expect_error(ValueError, rope_attention_logits, np.ones((1, 1, 8)), np.ones((1, 1, 8)), [0], [0], np.ones(4), pairing="unknown")
expect_error(ValueError, position_interpolated_frequencies, np.ones(4), 0.5)
expect_error(ValueError, ntk_scaled_base, 2, 10_000.0, 2.0)
expect_error(ValueError, yarn_scaled_frequencies, 64, 150_000.0, 32.0, 4096, beta_fast=1.0, beta_slow=32.0)
expect_error(ValueError, alibi_bias, np.array([-0.5]), np.arange(2), np.arange(2))
expect_error(TypeError, sinusoidal_encoding, np.array([0.0, 1.0]), 8)
expect_error(ValueError, apply_mrope_adjacent, np.ones((4, 8)), np.ones((4, 3), dtype=int), (1, 1, 1), np.ones(4))
expect_error(TypeError, apply_mrope_adjacent, np.ones((4, 8)), np.ones((4, 3), dtype=float), (2, 1, 1), np.ones(4))
print('Отрицательные тесты пройдены')

"""Behavioral checks embedded in the M4 notebook after its NumPy core."""
import numpy as np


def check_cached_positions():
    """Compare full causal attention with cached steps; require wrong IDs to fail."""
    rng = np.random.default_rng(404)
    q, k, v = [rng.standard_normal((7, 8)) for _ in range(3)]
    pos = np.arange(7, dtype=np.int64)
    inv = rope_inverse_frequencies(8, 10_000.0)
    qr = apply_rope_adjacent(q, pos, inv)
    kr = apply_rope_adjacent(k, pos, inv)

    def attend(queries, keys, values, allowed):
        scores = (queries @ keys.T) / np.sqrt(queries.shape[-1])
        scores = np.where(allowed, scores, -np.inf)
        weights = np.exp(scores - scores.max(axis=-1, keepdims=True))
        weights /= weights.sum(axis=-1, keepdims=True)
        return weights @ values

    full = attend(qr, kr, v, np.tril(np.ones((7, 7), dtype=bool)))
    cache_k, cache_v, steps, wrong = [], [], [], []
    for t in range(7):
        cache_k.append(apply_rope_adjacent(k[t:t+1], [t], inv))
        cache_v.append(v[t:t+1])
        keys, values = np.concatenate(cache_k), np.concatenate(cache_v)
        query = apply_rope_adjacent(q[t:t+1], [t], inv)
        steps.append(attend(query, keys, values, np.ones((1, t+1), dtype=bool)))
        reset = apply_rope_adjacent(q[t:t+1], [0], inv)
        wrong.append(attend(reset, keys, values, np.ones((1, t+1), dtype=bool)))
    cached = np.concatenate(steps)
    np.testing.assert_allclose(cached, full, rtol=1e-12, atol=1e-12)
    wrong_gap = float(np.max(np.abs(np.concatenate(wrong)-full)))
    assert wrong_gap > .01, 'The negative control must expose the reset-position bug'
    # Changing frequencies for Q alone leaves cached K in a different geometry.
    new_inv = inv / 4
    q_new = apply_rope_adjacent(q[-1:], [6], new_inv)
    new_keys = apply_rope_adjacent(k, pos, new_inv)
    mixed = attend(q_new, kr, v, np.ones((1, 7), dtype=bool))
    consistent = attend(q_new, new_keys, v, np.ones((1, 7), dtype=bool))
    spectrum_gap = float(np.max(np.abs(mixed-consistent)))
    assert spectrum_gap > .01
    return {'full_vs_cached_max_error': float(np.max(np.abs(cached-full))),
            'reset_query_position_gap': wrong_gap, 'mixed_spectrum_gap': spectrum_gap}


def check_axis_and_scaling_contracts():
    """Probe one M-RoPE axis and the full-versus-partial amplitude distinction."""
    inv = rope_inverse_frequencies(8, 10_000.0)
    x = np.ones((1, 8))
    first = apply_mrope_adjacent(x, [[0, 1, 2]], (2, 1, 1), inv)
    changed = apply_mrope_adjacent(x, [[0, 1, 3]], (2, 1, 1), inv)
    np.testing.assert_array_equal(first[:, :6], changed[:, :6])
    assert not np.allclose(first[:, 6:], changed[:, 6:])
    q = np.array([1., 2., 3., 4.]); k = np.array([2., -1., 1., 2.]); a = 1.3
    np.testing.assert_allclose((a*q) @ (a*k), a*a*(q@k), rtol=1e-14)
    partial_q, partial_k = q.copy(), k.copy()
    partial_q[2:] *= a; partial_k[2:] *= a
    expected = q[:2]@k[:2] + a*a*(q[2:]@k[2:])
    np.testing.assert_allclose(partial_q@partial_k, expected, rtol=1e-14)
    # Nonzero unrotated contribution is essential for the counterexample.
    k[0] = 3.; partial_k = k.copy(); partial_k[2:] *= a
    assert not np.isclose(partial_q@partial_k, a*a*(q@k))
    return {'only_width_pairs_changed': True, 'partial_rope_is_not_global_scaling': True}

print(check_cached_positions())
print(check_axis_and_scaling_contracts())

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_rope_explicit(
    x: np.ndarray,
    position_ids: np.ndarray,
    *,
    base: float = 10_000.0,
) -> np.ndarray:
    """Adjacent pairs; valid floating x[..., T, D], integer IDs[T], base > 1."""
    dim = x.shape[-1]
    inv_freq = base ** (-np.arange(0, dim, 2) / dim)
    angles = position_ids[:, None] * inv_freq[None, :]
    shape = (1,) * (x.ndim - 2) + angles.shape
    cos, sin = np.cos(angles).reshape(shape), np.sin(angles).reshape(shape)
    even, odd = x[..., 0::2], x[..., 1::2]
    out = np.empty_like(x, dtype=np.result_type(x.dtype, np.float64))
    out[..., 0::2] = even * cos - odd * sin
    out[..., 1::2] = even * sin + odd * cos
    return out

def torch_rope_explicit(
    x: torch.Tensor,
    position_ids: torch.Tensor,
    *,
    base: float = 10_000.0,
) -> torch.Tensor:
    """Adjacent pairs with at least float32 phase arithmetic; output keeps x.dtype."""
    dim = x.shape[-1]
    work_dtype = torch.float64 if x.dtype == torch.float64 else torch.float32
    inv_freq = base ** (-torch.arange(0, dim, 2, device=x.device, dtype=work_dtype) / dim)
    angles = position_ids.to(device=x.device, dtype=work_dtype)[:, None] * inv_freq[None, :]
    shape = (1,) * (x.ndim - 2) + tuple(angles.shape)
    cos, sin = angles.cos().reshape(shape), angles.sin().reshape(shape)
    even, odd = x[..., 0::2].to(work_dtype), x[..., 1::2].to(work_dtype)
    out = torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
    return out.to(x.dtype)

rng = np.random.default_rng(2)
x_np = rng.standard_normal((2, 5, 8))
pos_np = np.arange(5, dtype=np.int64) + 11
np_out = numpy_rope_explicit(x_np, pos_np)
x = torch.tensor(x_np, dtype=torch.float64, requires_grad=True)
t_out = torch_rope_explicit(x, torch.tensor(pos_np))
np.testing.assert_allclose(np_out, t_out.detach().numpy(), rtol=1e-12, atol=1e-12)
t_out.square().sum().backward()
# A rotation is orthogonal: d ||R x||^2 / dx = 2 x.
torch.testing.assert_close(x.grad, 2 * x.detach(), rtol=1e-12, atol=1e-12)
# Half-precision IDs would merge 2048 and 2049 before the rotation.
long_ids = torch.tensor([2048, 2049], dtype=torch.int64)
half_x = torch.ones((2, 8), dtype=torch.float16)
half_out = torch_rope_explicit(half_x, long_ids)
expected = numpy_rope_explicit(half_x.numpy(), long_ids.numpy())
np.testing.assert_allclose(half_out.numpy(), expected, rtol=1e-3, atol=1e-3)
assert not torch.equal(half_out[0], half_out[1])
assert long_ids.to(torch.float16)[0] == long_ids.to(torch.float16)[1]
print("B02 outputs, analytic gradient, and float16 positions: 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"воспроизводимый тестовый запуск; seed={SEED}")
DEMO_LANGUAGE = 'ru'

import importlib.util
import subprocess
import 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 hashlib
import json
from dataclasses import dataclass
from numbers import Integral, Real
from typing import Sequence


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

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


def _positive_even_int(name: str, value: object, *, minimum: int = 2) -> int:
    """Return an even integer not smaller than ``minimum``."""

    result = _positive_int(name, value)
    if result < minimum:
        raise ValueError(f"{name} must be at least {minimum}")
    if result % 2:
        raise ValueError(f"{name} must be even")
    return result


def _finite_real(name: str, value: object) -> float:
    """Return a finite real scalar and reject booleans and complex values."""

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


def _real_greater_than(name: str, value: object, lower: float) -> float:
    """Return a finite scalar strictly greater than ``lower``."""

    result = _finite_real(name, value)
    if result <= lower:
        raise ValueError(f"{name} must be greater than {lower}")
    return result


def _finite_array(
    name: str,
    value: object,
    *,
    ndim: int | None = None,
    minimum_ndim: int | None = None,
) -> np.ndarray:
    """Convert to float64 and enforce rank, non-emptiness, and finite values."""

    try:
        if np.iscomplexobj(value):
            raise TypeError(f"{name} must be real")
        array = np.asarray(value, dtype=np.float64)
    except (TypeError, ValueError) as exc:
        raise TypeError(f"{name} must be numeric") from exc
    if ndim is not None and array.ndim != ndim:
        raise ValueError(f"{name} must have rank {ndim}, got shape {array.shape}")
    if minimum_ndim is not None and array.ndim < minimum_ndim:
        raise ValueError(f"{name} must have rank at least {minimum_ndim}, got shape {array.shape}")
    if 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 _position_vector(name: str, value: object) -> np.ndarray:
    """Validate a non-empty one-dimensional vector of non-negative integer positions."""

    raw = np.asarray(value)
    if raw.ndim != 1 or raw.size == 0:
        raise ValueError(f"{name} must be a non-empty one-dimensional vector")
    if raw.dtype.kind not in "iu":
        raise TypeError(f"{name} must contain integer positions")
    result = raw.astype(np.int64, copy=False)
    if np.any(result < 0):
        raise ValueError(f"{name} must contain non-negative positions")
    return result


def _inverse_frequency_vector(value: object, *, expected_pairs: int | None = None) -> np.ndarray:
    """Validate a positive one-dimensional inverse-frequency vector."""

    inv = _finite_array("inverse_frequencies", value, ndim=1)
    if np.any(inv <= 0.0):
        raise ValueError("inverse_frequencies must be strictly positive")
    if expected_pairs is not None and inv.size != expected_pairs:
        raise ValueError(
            f"inverse_frequencies must have {expected_pairs} entries, got {inv.size}"
        )
    return inv


def _sequence_tensor(name: str, value: object) -> np.ndarray:
    """Validate a finite sequence tensor with shape ``[..., sequence, head_dim]``."""

    tensor = _finite_array(name, value, minimum_ndim=2)
    _positive_even_int(f"{name}.shape[-1]", tensor.shape[-1])
    if tensor.shape[-2] <= 0:
        raise ValueError(f"{name} must have a non-empty sequence axis")
    return tensor

def rope_inverse_frequencies(head_dim: int, base: Real = 10_000.0) -> np.ndarray:
    """Return ``base**(-2i/head_dim)`` for the ``head_dim/2`` RoPE pairs."""

    dim = _positive_even_int("head_dim", head_dim)
    rope_base = _real_greater_than("base", base, 1.0)
    pair_indices = np.arange(dim // 2, dtype=np.float64)
    inv = rope_base ** (-2.0 * pair_indices / dim)
    if not np.all(np.isfinite(inv)) or np.any(inv <= 0.0):
        raise FloatingPointError("inverse-frequency construction produced invalid values")
    return inv


def sinusoidal_encoding(
    positions: object,
    model_dim: int,
    *,
    base: Real = 10_000.0,
) -> np.ndarray:
    """Return the classic sinusoidal encoding with shape ``[positions, model_dim]``."""

    pos = _position_vector("positions", positions)
    dim = _positive_even_int("model_dim", model_dim)
    inv = rope_inverse_frequencies(dim, base)
    angles = pos.astype(np.float64)[:, None] * inv[None, :]
    result = np.empty((pos.size, dim), dtype=np.float64)
    result[:, 0::2] = np.sin(angles)
    result[:, 1::2] = np.cos(angles)
    if not np.all(np.isfinite(result)):
        raise FloatingPointError("sinusoidal encoding produced non-finite values")
    return result


def wavelengths_from_inverse_frequencies(inverse_frequencies: object) -> np.ndarray:
    """Convert radians-per-position frequencies to tokens per full ``2π`` turn."""

    inv = _inverse_frequency_vector(inverse_frequencies)
    return (2.0 * math.pi) / inv


def rotations_within_context(
    inverse_frequencies: object,
    context_length: int,
) -> np.ndarray:
    """Return turns (including fractions) across a span L = context_length.

    This is the YaRN L*theta/(2*pi) diagnostic. Positions 0..L-1 have a maximum
    endpoint separation of L-1, which is a different convention.
    """

    inv = _inverse_frequency_vector(inverse_frequencies)
    length = _positive_int("context_length", context_length)
    return length * inv / (2.0 * math.pi)


def full_turn_count(inverse_frequencies: object, context_length: int) -> int:
    """Count frequency pairs that complete at least one full turn in the context."""

    return int(np.count_nonzero(rotations_within_context(inverse_frequencies, context_length) >= 1.0))

def rotate_half_split(x: object) -> np.ndarray:
    """Apply the split-half quarter turn used by HF/Llama-style RoPE implementations."""

    values = _sequence_tensor("x", x)
    half = values.shape[-1] // 2
    return np.concatenate((-values[..., half:], values[..., :half]), axis=-1)


def _angles_for_sequence(
    x: np.ndarray,
    positions: np.ndarray,
    inverse_frequencies: np.ndarray,
) -> np.ndarray:
    """Build a broadcastable angle tensor for ``[..., T, D]`` sequence tensors."""

    if positions.size != x.shape[-2]:
        raise ValueError(
            f"positions length {positions.size} does not match sequence axis {x.shape[-2]}"
        )
    if inverse_frequencies.size != x.shape[-1] // 2:
        raise ValueError("inverse_frequencies do not match head_dim")
    base_angles = positions.astype(np.float64)[:, None] * inverse_frequencies[None, :]
    return base_angles.reshape((1,) * (x.ndim - 2) + base_angles.shape)


def apply_rope_adjacent(
    x: object,
    positions: object,
    inverse_frequencies: object,
) -> np.ndarray:
    """Apply RoPE to adjacent pairs ``(x_0,x_1), (x_2,x_3), ...``."""

    values = _sequence_tensor("x", x)
    pos = _position_vector("positions", positions)
    inv = _inverse_frequency_vector(
        inverse_frequencies,
        expected_pairs=values.shape[-1] // 2,
    )
    angles = _angles_for_sequence(values, pos, inv)
    cosine = np.cos(angles)
    sine = np.sin(angles)
    first = values[..., 0::2]
    second = values[..., 1::2]
    result = np.empty_like(values, dtype=np.float64)
    result[..., 0::2] = cosine * first - sine * second
    result[..., 1::2] = sine * first + cosine * second
    if not np.all(np.isfinite(result)):
        raise FloatingPointError("adjacent-pair RoPE produced non-finite values")
    return result


def apply_rope_split_half(
    x: object,
    positions: object,
    inverse_frequencies: object,
) -> np.ndarray:
    """Apply split-half RoPE to pairs ``(x_j, x_{j+D/2})``."""

    values = _sequence_tensor("x", x)
    pos = _position_vector("positions", positions)
    inv = _inverse_frequency_vector(
        inverse_frequencies,
        expected_pairs=values.shape[-1] // 2,
    )
    angles = _angles_for_sequence(values, pos, inv)
    repeated = np.concatenate((angles, angles), axis=-1)
    result = values * np.cos(repeated) + rotate_half_split(values) * np.sin(repeated)
    if not np.all(np.isfinite(result)):
        raise FloatingPointError("split-half RoPE produced non-finite values")
    return result


def apply_rope_complex_adjacent(
    x: object,
    positions: object,
    inverse_frequencies: object,
) -> np.ndarray:
    """Apply adjacent-pair RoPE through the equivalent complex multiplication."""

    values = _sequence_tensor("x", x)
    pos = _position_vector("positions", positions)
    inv = _inverse_frequency_vector(
        inverse_frequencies,
        expected_pairs=values.shape[-1] // 2,
    )
    angles = _angles_for_sequence(values, pos, inv)
    complex_values = values[..., 0::2] + 1j * values[..., 1::2]
    rotated = complex_values * np.exp(1j * angles)
    result = np.empty_like(values, dtype=np.float64)
    result[..., 0::2] = rotated.real
    result[..., 1::2] = rotated.imag
    if not np.all(np.isfinite(result)):
        raise FloatingPointError("complex RoPE produced non-finite values")
    return result


def adjacent_to_split_permutation(head_dim: int) -> np.ndarray:
    """Return the coordinate permutation mapping adjacent pairs to split-half pairs."""

    dim = _positive_even_int("head_dim", head_dim)
    return np.concatenate((np.arange(0, dim, 2), np.arange(1, dim, 2)))


def rope_attention_logits(
    queries: object,
    keys: object,
    query_positions: object,
    key_positions: object,
    inverse_frequencies: object,
    *,
    pairing: str = "split_half",
) -> np.ndarray:
    """Return unscaled rotated dot products ``[..., query, key]``.

    Divide by sqrt(head_dim) before the standard attention softmax. The historical
    function name and fixture key 'logits' are retained for file compatibility.
    """

    q = _sequence_tensor("queries", queries)
    k = _sequence_tensor("keys", keys)
    if q.ndim != k.ndim or q.shape[:-2] != k.shape[:-2]:
        raise ValueError("queries and keys must have identical batch dimensions")
    if q.shape[-1] != k.shape[-1]:
        raise ValueError("queries and keys must have the same head_dim")
    if pairing == "split_half":
        apply = apply_rope_split_half
    elif pairing == "adjacent":
        apply = apply_rope_adjacent
    else:
        raise ValueError("pairing must be 'split_half' or 'adjacent'")
    q_rot = apply(q, query_positions, inverse_frequencies)
    k_rot = apply(k, key_positions, inverse_frequencies)
    return np.matmul(q_rot, np.swapaxes(k_rot, -1, -2))

def position_interpolated_frequencies(
    inverse_frequencies: object,
    factor: Real,
) -> np.ndarray:
    """Apply position interpolation by dividing every frequency by ``factor``."""

    inv = _inverse_frequency_vector(inverse_frequencies)
    scale = _finite_real("factor", factor)
    if scale < 1.0:
        raise ValueError("factor must be at least 1")
    return inv / scale


def ntk_scaled_base(head_dim: int, base: Real, factor: Real) -> float:
    """Return the static NTK-aware base ``base * factor**(d/(d-2))``."""

    dim = _positive_even_int("head_dim", head_dim, minimum=4)
    rope_base = _real_greater_than("base", base, 1.0)
    scale = _finite_real("factor", factor)
    if scale < 1.0:
        raise ValueError("factor must be at least 1")
    result = rope_base * scale ** (dim / (dim - 2))
    if not math.isfinite(result):
        raise FloatingPointError("NTK base is non-finite")
    return result


def ntk_scaled_frequencies(head_dim: int, base: Real, factor: Real) -> np.ndarray:
    """Return inverse frequencies generated from the static NTK-aware base."""

    return rope_inverse_frequencies(head_dim, ntk_scaled_base(head_dim, base, factor))


@dataclass(frozen=True)
class YarnSpectrum:
    """Piecewise YaRN spectrum plus diagnostic boundaries and scaling metadata."""

    inverse_frequencies: np.ndarray
    original_inverse_frequencies: np.ndarray
    ramp: np.ndarray
    low: float
    high: float
    truncate: bool
    kept_pairs: int
    transition_pairs: int
    interpolated_pairs: int
    attention_factor: float
    score_factor_if_applied_to_q_and_k: float


def yarn_scaled_frequencies(
    head_dim: int,
    base: Real,
    factor: Real,
    original_context: int,
    *,
    beta_fast: Real = 32.0,
    beta_slow: Real = 1.0,
    truncate: bool = False,
) -> YarnSpectrum:
    """Construct the lecture's piecewise YaRN frequency spectrum.

    ``ramp == 0`` keeps a frequency, ``ramp == 1`` divides it by ``factor``.
    With ``truncate=False`` this follows the continuous correction bounds used by the
    original gpt-oss implementation. ``truncate=True`` reproduces the integer-boundary
    option supported by the generic Hugging Face YaRN utility.
    """

    dim = _positive_even_int("head_dim", head_dim, minimum=4)
    rope_base = _real_greater_than("base", base, 1.0)
    scale = _finite_real("factor", factor)
    if scale < 1.0:
        raise ValueError("factor must be at least 1")
    context = _positive_int("original_context", original_context)
    fast = _real_greater_than("beta_fast", beta_fast, 0.0)
    slow = _real_greater_than("beta_slow", beta_slow, 0.0)
    if fast <= slow:
        raise ValueError("beta_fast must be greater than beta_slow")
    if not isinstance(truncate, bool):
        raise TypeError("truncate must be a bool")

    original = rope_inverse_frequencies(dim, rope_base)
    pair_count = dim // 2

    def correction_dimension(rotations: float) -> float:
        """Map a target turn count to the corresponding continuous pair index."""

        return dim * math.log(context / (rotations * 2.0 * math.pi)) / (2.0 * math.log(rope_base))

    low_raw = correction_dimension(fast)
    high_raw = correction_dimension(slow)
    low = math.floor(low_raw) if truncate else low_raw
    high = math.ceil(high_raw) if truncate else high_raw
    low = max(0.0, min(float(pair_count - 1), float(low)))
    high = max(0.0, min(float(pair_count - 1), float(high)))
    if high <= low:
        raise ValueError("YaRN correction range is empty for this configuration")
    pair_indices = np.arange(pair_count, dtype=np.float64)
    ramp = np.clip((pair_indices - low) / (high - low), 0.0, 1.0)
    scaled = original * (1.0 - ramp) + (original / scale) * ramp
    attention_factor = 0.1 * math.log(scale) + 1.0 if scale > 1.0 else 1.0
    score_factor_if_applied_to_q_and_k = attention_factor**2
    kept = int(np.count_nonzero(ramp == 0.0))
    interpolated = int(np.count_nonzero(ramp == 1.0))
    transition = pair_count - kept - interpolated
    return YarnSpectrum(
        inverse_frequencies=scaled,
        original_inverse_frequencies=original,
        ramp=ramp,
        low=low,
        high=high,
        truncate=truncate,
        kept_pairs=kept,
        transition_pairs=transition,
        interpolated_pairs=interpolated,
        attention_factor=attention_factor,
        score_factor_if_applied_to_q_and_k=score_factor_if_applied_to_q_and_k,
    )

def alibi_bias(
    slopes: object,
    query_positions: object,
    key_positions: object,
    *,
    causal: bool = True,
) -> np.ndarray:
    """Return an ALiBi tensor with shape ``[heads, query, key]``.

    In causal mode, future keys receive ``-inf`` and allowed keys receive
    ``-slope * (query_position - key_position)``. In bidirectional mode the absolute
    distance is used.
    """

    if not isinstance(causal, bool):
        raise TypeError("causal must be a bool")
    slope_values = _finite_array("slopes", slopes, ndim=1)
    if np.any(slope_values <= 0.0):
        raise ValueError("slopes must be strictly positive")
    q_pos = _position_vector("query_positions", query_positions)
    k_pos = _position_vector("key_positions", key_positions)
    delta = q_pos[:, None] - k_pos[None, :]
    if causal:
        allowed = delta >= 0
        base = -np.abs(delta).astype(np.float64)
        bias = slope_values[:, None, None] * base[None, :, :]
        bias = np.where(allowed[None, :, :], bias, -np.inf)
    else:
        bias = -slope_values[:, None, None] * np.abs(delta)[None, :, :]
    return bias


def apply_mrope_adjacent(
    x: object,
    coordinates: object,
    axis_sections: Sequence[int],
    inverse_frequencies: object,
) -> np.ndarray:
    """Apply adjacent-pair multi-axis RoPE.

    ``coordinates`` has shape ``[sequence, axes]``. ``axis_sections`` gives the number of
    frequency pairs assigned to each axis and must sum to ``head_dim / 2``.
    """

    values = _sequence_tensor("x", x)
    raw_coordinates = np.asarray(coordinates)
    if raw_coordinates.ndim != 2 or raw_coordinates.shape[0] != values.shape[-2]:
        raise ValueError("coordinates must have shape [sequence, axes]")
    if raw_coordinates.dtype.kind not in "iu":
        raise TypeError("coordinates must contain integer positions")
    coords = raw_coordinates.astype(np.int64, copy=False)
    if np.any(coords < 0):
        raise ValueError("coordinates must be non-negative")

    if isinstance(axis_sections, (str, bytes)):
        raise TypeError("axis_sections must be a sequence of integers")
    sections: list[int] = []
    for index, value in enumerate(axis_sections):
        sections.append(_positive_int(f"axis_sections[{index}]", value))
    if len(sections) != coords.shape[1]:
        raise ValueError("axis_sections length must match the coordinate axis count")
    pair_count = values.shape[-1] // 2
    if sum(sections) != pair_count:
        raise ValueError("axis_sections must sum to head_dim / 2")

    inv = _inverse_frequency_vector(inverse_frequencies, expected_pairs=pair_count)
    axis_index = np.repeat(np.arange(len(sections)), sections)
    base_angles = coords[:, axis_index].astype(np.float64) * inv[None, :]
    angles = base_angles.reshape((1,) * (values.ndim - 2) + base_angles.shape)
    cosine = np.cos(angles)
    sine = np.sin(angles)
    first = values[..., 0::2]
    second = values[..., 1::2]
    result = np.empty_like(values, dtype=np.float64)
    result[..., 0::2] = cosine * first - sine * second
    result[..., 1::2] = sine * first + cosine * second
    if not np.all(np.isfinite(result)):
        raise FloatingPointError("M-RoPE produced non-finite values")
    return result
INK = '#2E3440'
MUTED = '#4C566A'
BORDER = '#D8DEE9'
BLUE = '#5E81AC'
ORANGE = '#D08770'
PURPLE = '#B48EAD'
GREEN = '#A3BE8C'
PALE = {'blue': '#E8EEF5', 'green': '#EFF4EA', 'purple': '#F3EEF3', 'orange': '#FAF0E8', 'yellow': '#FAF5E3', 'gray': '#ECEFF4'}

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

def scaling_plot(reference, lang='en'):
    plot_style()
    ru = lang == 'ru'; sc = reference['scaling']; yr = sc['yarn']; x = np.arange(len(sc['raw']))
    fig, (ax, ramp_ax) = plt.subplots(2, 1, figsize=(11, 7), sharex=True,
                                    height_ratios=[3, 1], layout='constrained')
    series = [('Исходные' if ru else 'Original', sc['raw'], MUTED, '-'),
              ('PI', sc['position_interpolation'], PURPLE, '--'),
              ('Статическая NTK-aware' if ru else 'Static NTK-aware', sc['ntk'], ORANGE, '-.'),
              ('YaRN', yr['inverse_frequencies'], BLUE, '-')]
    for name, values, color, style in series:
        ax.plot(x, values, linestyle=style, color=color, linewidth=2.2, label=name)
    ax.set_yscale('log'); ax.grid(axis='y')
    ax.set_ylabel('Частота, рад/шаг (лог. шкала)' if ru else 'Frequency, rad/step (log scale)')
    ax.set_title('Как меняется спектр · s=32' if ru else 'How the spectrum changes · s=32', pad=16)
    ax.legend(loc='lower left', ncol=2)
    ramp_ax.plot(x, yr['ramp'], color=BLUE, marker='o', markersize=3)
    for value in (yr['low'], yr['high']):
        ramp_ax.axvline(value, color=MUTED, linestyle=':', linewidth=1.4)
    ramp_ax.set(xlabel='Индекс пары i' if ru else 'Pair index i',
                ylabel='Вес PI' if ru else 'PI weight', ylim=(-.08,1.08), yticks=[0,.5,1])
    ramp_ax.grid(axis='y')
    return fig

reference = json.loads('{\n  "alibi": {\n    "allowed_bias_with_null_for_future": [\n      [\n        [\n          -0.0,\n          null,\n          null,\n          null\n        ],\n        [\n          -1.0,\n          -0.0,\n          null,\n          null\n        ],\n        [\n          -2.0,\n          -1.0,\n          -0.0,\n          null\n        ],\n        [\n          -3.0,\n          -2.0,\n          -1.0,\n          -0.0\n        ]\n      ],\n      [\n        [\n          -0.0,\n          null,\n          null,\n          null\n        ],\n        [\n          -0.5,\n          -0.0,\n          null,\n          null\n        ],\n        [\n          -1.0,\n          -0.5,\n          -0.0,\n          null\n        ],\n        [\n          -1.5,\n          -1.0,\n          -0.5,\n          -0.0\n        ]\n      ]\n    ],\n    "positions": [\n      0,\n      1,\n      2,\n      3\n    ],\n    "slopes": [\n      1.0,\n      0.5\n    ]\n  },\n  "meta": {\n    "description": "Frozen numerical contract independent of lecture prose and figures",\n    "module": "module-04",\n    "version": 2\n  },\n  "mrope": {\n    "axis_sections": [\n      2,\n      1,\n      1\n    ],\n    "base": 10000.0,\n    "head_dim": 8,\n    "image_coordinates": [\n      [\n        0,\n        0,\n        0\n      ],\n      [\n        0,\n        0,\n        1\n      ],\n      [\n        0,\n        1,\n        0\n      ],\n      [\n        0,\n        1,\n        1\n      ]\n    ],\n    "image_input": [\n      [\n        0.0,\n        0.1,\n        0.2,\n        0.3,\n        0.4,\n        0.5,\n        0.6,\n        0.7\n      ],\n      [\n        0.8,\n        0.9,\n        1.0,\n        1.1,\n        1.2,\n        1.3,\n        1.4,\n        1.5\n      ],\n      [\n        1.6,\n        1.7,\n        1.8,\n        1.9,\n        2.0,\n        2.1,\n        2.2,\n        2.3\n      ],\n      [\n        2.4,\n        2.5,\n        2.6,\n        2.7,\n        2.8,\n        2.9,\n        3.0,\n        3.1\n      ]\n    ],\n    "image_output": [\n      [\n        0.0,\n        0.1,\n        0.2,\n        0.3,\n        0.4,\n        0.5,\n        0.6,\n        0.7\n      ],\n      [\n        0.8,\n        0.9,\n        1.0,\n        1.1,\n        1.2,\n        1.3,\n        1.3984993002500583,\n        1.5013992497667292\n      ],\n      [\n        1.6,\n        1.7,\n        1.8,\n        1.9,\n        1.9789003508315806,\n        2.1198946675433303,\n        2.2,\n        2.3\n      ],\n      [\n        2.4,\n        2.5,\n        2.6,\n        2.7,\n        2.7708604844975793,\n        2.9278545345439957,\n        2.996898500516792,\n        3.1029984495001295\n      ]\n    ],\n    "text_coordinates": [\n      [\n        4,\n        4,\n        4\n      ],\n      [\n        5,\n        5,\n        5\n      ],\n      [\n        6,\n        6,\n        6\n      ]\n    ],\n    "text_input": [\n      [\n        0.0,\n        0.1,\n        0.2,\n        0.3,\n        0.4,\n        0.5,\n        0.6,\n        0.7\n      ],\n      [\n        0.8,\n        0.9,\n        1.0,\n        1.1,\n        1.2,\n        1.3,\n        1.4,\n        1.5\n      ],\n      [\n        1.6,\n        1.7,\n        1.8,\n        1.9,\n        2.0,\n        2.1,\n        2.2,\n        2.3\n      ]\n    ],\n    "text_matches_standard_1d": true,\n    "text_output": [\n      [\n        0.07568024953079283,\n        -0.06536436208636119,\n        0.0673866961079819,\n        0.3542019666625956,\n        0.37968537557107407,\n        0.5155957870051426,\n        0.5971952074730607,\n        0.7023943936074717\n      ],\n      [\n        1.0899615955674056,\n        -0.5118434528136071,\n        0.3502144694257494,\n        1.444766356683613,\n        1.1335273924220777,\n        1.3583503416382703,\n        1.3924825312864193,\n        1.5069812208724322\n      ],\n      [\n        2.0112788055787596,\n        1.185224690187341,\n        0.4127834073868535,\n        2.5844941204394525,\n        1.8704766662635748,\n        2.216149146822818,\n        2.1861604829186514,\n        2.3131585209243424\n      ]\n    ],\n    "text_positions": [\n      4,\n      5,\n      6\n    ],\n    "text_standard_rope_output": [\n      [\n        0.07568024953079283,\n        -0.06536436208636119,\n        0.0673866961079819,\n        0.3542019666625956,\n        0.37968537557107407,\n        0.5155957870051426,\n        0.5971952074730607,\n        0.7023943936074717\n      ],\n      [\n        1.0899615955674056,\n        -0.5118434528136071,\n        0.3502144694257494,\n        1.444766356683613,\n        1.1335273924220777,\n        1.3583503416382703,\n        1.3924825312864193,\n        1.5069812208724322\n      ],\n      [\n        2.0112788055787596,\n        1.185224690187341,\n        0.4127834073868535,\n        2.5844941204394525,\n        1.8704766662635748,\n        2.216149146822818,\n        2.1861604829186514,\n        2.3131585209243424\n      ]\n    ]\n  },\n  "scaling": {\n    "base": 150000.0,\n    "beta_fast": 32.0,\n    "beta_slow": 1.0,\n    "extended_context": 131072,\n    "factor": 32.0,\n    "head_dim": 64,\n    "ntk": [\n      1.0,\n      0.6161605983892874,\n      0.37965388300744474,\n      0.23392776373468369,\n      0.14413707088263056,\n      0.08881158384512078,\n      0.05472219864590999,\n      0.033717662662841354,\n      0.020775495202624464,\n      0.01280104155588286,\n      0.007887497425078919,\n      0.00485996513323059,\n      0.0029945190246424332,\n      0.001845104634111787,\n      0.0011368807754451658,\n      0.0007005011388955706,\n      0.00043162120091427207,\n      0.00026594797743284077,\n      0.00016386666491543985,\n      0.00010096818231035427,\n      6.221261563062656e-05,\n      3.83329624743296e-05,\n      2.3619261096217023e-05,\n      1.4553258050557898e-05,\n      8.967144188945469e-06,\n      5.525200929303662e-06,\n      3.4044111108207913e-06,\n      2.0976639872064776e-06,\n      1.2924978975768017e-06,\n      7.963862779878181e-07,\n      4.907018455939914e-07,\n      3.0235114281192144e-07\n    ],\n    "ntk_base": 5367776.967415848,\n    "original_context": 4096,\n    "position_interpolation": [\n      0.03125,\n      0.0215326345590051,\n      0.014836939233653148,\n      0.01022330849575739,\n      0.007044312506339714,\n      0.0048538434213901555,\n      0.0033445131711816074,\n      0.0023045177565706506,\n      0.001587914837983171,\n      0.0010941436773476293,\n      0.0007539134706999197,\n      0.0005194797841181623,\n      0.00035794458727063204,\n      0.00024663967968231775,\n      0.00016994566689118156,\n      0.0001171000940817352,\n      8.068715304598785e-05,\n      5.55970233646647e-05,\n      3.8308812373753384e-05,\n      2.639646903147278e-05,\n      1.818833668168956e-05,\n      1.2532569824095001e-05,\n      8.63549587543848e-06,\n      5.950239261491605e-06,\n      4.09997848180298e-06,\n      2.8250668271503398e-06,\n      1.9465962110014815e-06,\n      1.3412910350540445e-06,\n      9.242089502428302e-07,\n      6.36820914615699e-07,\n      4.387978250864396e-07,\n      3.0235114281192144e-07\n    ],\n    "raw": [\n      1.0,\n      0.6890443058881632,\n      0.4747820554769007,\n      0.3271458718642365,\n      0.22541800020287084,\n      0.15532298948448497,\n      0.10702442147781144,\n      0.07374456821026082,\n      0.050813274815461475,\n      0.03501259767512414,\n      0.02412523106239743,\n      0.016623353091781194,\n      0.011454226792660225,\n      0.007892469749834168,\n      0.00543826134051781,\n      0.0037472030106155264,\n      0.0025819888974716113,\n      0.0017791047476692703,\n      0.0012258819959601083,\n      0.0008446870090071289,\n      0.0005820267738140659,\n      0.00040104223437104005,\n      0.00027633586801403137,\n      0.00019040765636773135,\n      0.00013119931141769537,\n      9.040213846881087e-05,\n      6.229107875204741e-05,\n      4.292131312172942e-05,\n      2.9574686407770566e-05,\n      2.0378269267702367e-05,\n      1.4041530402766067e-05,\n      9.675236569981486e-06\n    ],\n    "yarn": {\n      "attention_factor": 1.3465735902799727,\n      "hf_truncated_high": 18,\n      "hf_truncated_low": 8,\n      "high": 17.39802450158856,\n      "interpolated_pairs": 14,\n      "inverse_frequencies": [\n        1.0,\n        0.6890443058881632,\n        0.4747820554769007,\n        0.3271458718642365,\n        0.22541800020287084,\n        0.15532298948448497,\n        0.10702442147781144,\n        0.07374456821026082,\n        0.050813274815461475,\n        0.03170569618466377,\n        0.01933500112654036,\n        0.011592049256286924,\n        0.006794959489732219,\n        0.0038603593171920685,\n        0.0020937923789696887,\n        0.0010526021013863357,\n        0.00045648391922324086,\n        0.0001293187012450632,\n        3.8308812373753384e-05,\n        2.639646903147278e-05,\n        1.818833668168956e-05,\n        1.2532569824095001e-05,\n        8.63549587543848e-06,\n        5.950239261491605e-06,\n        4.09997848180298e-06,\n        2.8250668271503398e-06,\n        1.9465962110014815e-06,\n        1.3412910350540445e-06,\n        9.242089502428302e-07,\n        6.36820914615699e-07,\n        4.387978250864396e-07,\n        3.0235114281192144e-07\n      ],\n      "kept_pairs": 9,\n      "low": 8.092779115512402,\n      "ramp": [\n        0.0,\n        0.0,\n        0.0,\n        0.0,\n        0.0,\n        0.0,\n        0.0,\n        0.0,\n        0.0,\n        0.09749564324709936,\n        0.20496191184183662,\n        0.3124281804365739,\n        0.4198944490313112,\n        0.5273607176260484,\n        0.6348269862207857,\n        0.742293254815523,\n        0.8497595234102603,\n        0.9572257920049976,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        1.0,\n        1.0\n      ],\n      "score_factor_if_applied_to_q_and_k": 1.8132604340394958,\n      "transition_pairs": 9,\n      "truncate": false\n    }\n  },\n  "sinusoidal": {\n    "base": 10000.0,\n    "encoding": [\n      [\n        0.0,\n        1.0,\n        0.0,\n        1.0,\n        0.0,\n        1.0,\n        0.0,\n        1.0\n      ],\n      [\n        0.8414709848078965,\n        0.5403023058681398,\n        0.09983341664682815,\n        0.9950041652780258,\n        0.009999833334166664,\n        0.9999500004166653,\n        0.0009999998333333417,\n        0.9999995000000417\n      ],\n      [\n        0.9092974268256817,\n        -0.4161468365471424,\n        0.19866933079506122,\n        0.9800665778412416,\n        0.01999866669333308,\n        0.9998000066665778,\n        0.0019999986666669333,\n        0.9999980000006666\n      ],\n      [\n        0.1411200080598672,\n        -0.9899924966004454,\n        0.2955202066613396,\n        0.955336489125606,\n        0.02999550020249566,\n        0.9995500337489875,\n        0.002999995500002025,\n        0.999995500003375\n      ]\n    ],\n    "model_dim": 8,\n    "positions": [\n      0,\n      1,\n      2,\n      3\n    ]\n  },\n  "spectrum": {\n    "bases": {\n      "10000": {\n        "full_turn_counts": {\n          "131072": 32,\n          "4096": 23\n        },\n        "inverse_frequencies": [\n          1.0,\n          0.7498942093324558,\n          0.5623413251903491,\n          0.4216965034285822,\n          0.31622776601683794,\n          0.23713737056616552,\n          0.1778279410038923,\n          0.1333521432163324,\n          0.1,\n          0.07498942093324558,\n          0.056234132519034905,\n          0.042169650342858224,\n          0.03162277660168379,\n          0.023713737056616554,\n          0.01778279410038923,\n          0.01333521432163324,\n          0.01,\n          0.007498942093324558,\n          0.005623413251903491,\n          0.004216965034285822,\n          0.0031622776601683794,\n          0.0023713737056616554,\n          0.0017782794100389228,\n          0.001333521432163324,\n          0.001,\n          0.0007498942093324559,\n          0.0005623413251903491,\n          0.00042169650342858224,\n          0.00031622776601683794,\n          0.0002371373705661655,\n          0.00017782794100389227,\n          0.0001333521432163324\n        ],\n        "maximum_wavelength_tokens": 47117.2427801674,\n        "minimum_wavelength_tokens": 6.283185307179586,\n        "wavelengths": [\n          6.283185307179586,\n          8.378762269377678,\n          11.173259061216543,\n          14.899780425245321,\n          19.869176531592203,\n          26.49597274431474,\n          35.33294752055899,\n          47.117242780167395,\n          62.83185307179586,\n          83.78762269377677,\n          111.73259061216544,\n          148.99780425245322,\n          198.69176531592203,\n          264.95972744314736,\n          353.32947520558986,\n          471.1724278016739,\n          628.3185307179587,\n          837.8762269377677,\n          1117.3259061216543,\n          1489.9780425245322,\n          1986.9176531592202,\n          2649.597274431474,\n          3533.294752055899,\n          4711.724278016739,\n          6283.185307179586,\n          8378.762269377676,\n          11173.259061216542,\n          14899.78042524532,\n          19869.176531592202,\n          26495.972744314742,\n          35332.94752055899,\n          47117.2427801674\n        ]\n      },\n      "150000": {\n        "full_turn_counts": {\n          "131072": 27,\n          "4096": 18\n        },\n        "inverse_frequencies": [\n          1.0,\n          0.6890443058881632,\n          0.4747820554769007,\n          0.3271458718642365,\n          0.22541800020287084,\n          0.15532298948448497,\n          0.10702442147781144,\n          0.07374456821026082,\n          0.050813274815461475,\n          0.03501259767512414,\n          0.02412523106239743,\n          0.016623353091781194,\n          0.011454226792660225,\n          0.007892469749834168,\n          0.00543826134051781,\n          0.0037472030106155264,\n          0.0025819888974716113,\n          0.0017791047476692703,\n          0.0012258819959601083,\n          0.0008446870090071289,\n          0.0005820267738140659,\n          0.00040104223437104005,\n          0.00027633586801403137,\n          0.00019040765636773135,\n          0.00013119931141769537,\n          9.040213846881087e-05,\n          6.229107875204741e-05,\n          4.292131312172942e-05,\n          2.9574686407770566e-05,\n          2.0378269267702367e-05,\n          1.4041530402766067e-05,\n          9.675236569981486e-06\n        ],\n        "maximum_wavelength_tokens": 649408.9588128396,\n        "minimum_wavelength_tokens": 6.283185307179586,\n        "wavelengths": [\n          6.283185307179586,\n          9.118695638999144,\n          13.233830627546281,\n          19.206066307286523,\n          27.873485265262175,\n          40.452384595695705,\n          58.70795861748462,\n          85.20200822472704,\n          123.65243787176136,\n          179.4549883296344,\n          260.4404198628719,\n          377.9734011837886,\n          548.5473110420513,\n          796.0987506238595,\n          1155.3665618028806,\n          1676.7667215733509,\n          2433.467205584167,\n          3531.6556348977883,\n          5125.4405627016395,\n          7438.477495427609,\n          10795.354423311823,\n          15667.141185350692,\n          22737.494601535214,\n          32998.59589177952,\n          47890.3832595279,\n          69502.61811945203,\n          100868.14088081702,\n          146388.46880941882,\n          212451.46002727243,\n          308327.72031027393,\n          447471.54526274785,\n          649408.9588128396\n        ]\n      },\n      "500000": {\n        "full_turn_counts": {\n          "131072": 25,\n          "4096": 16\n        },\n        "inverse_frequencies": [\n          1.0,\n          0.6636012376960885,\n          0.44036660267178046,\n          0.2922278225730151,\n          0.19392274474868576,\n          0.12868737343265052,\n          0.08539710028576561,\n          0.05666962144529105,\n          0.03760603093086393,\n          0.024955408670558694,\n          0.016560440080994446,\n          0.010989528534539826,\n          0.007292664737217109,\n          0.004839421345719893,\n          0.003211445994752591,\n          0.0021311195369119653,\n          0.001414213562373095,\n          0.0009384738703573802,\n          0.000622772421914596,\n          0.0004132725499855165,\n          0.0002742481756762073,\n          0.00018199142881462546,\n          0.00012076973741146504,\n          8.01429472224798e-05,\n          5.318295896944988e-05,\n          3.529227739646723e-05,\n          2.341999896140934e-05,\n          1.5541540297632344e-05,\n          1.031338537721246e-05,\n          6.8439753011549275e-06,\n          4.5416704806078695e-06,\n          3.013858152139171e-06\n        ],\n        "maximum_wavelength_tokens": 2084764.7732591254,\n        "minimum_wavelength_tokens": 6.283185307179586,\n        "wavelengths": [\n          6.283185307179586,\n          9.468314629722128,\n          14.268078616903308,\n          21.500982527458316,\n          32.40045573469106,\n          48.8251888245115,\n          73.57609668424418,\n          110.87395939719457,\n          167.07919319459117,\n          251.77649423117703,\n          379.40931982782695,\n          571.7429357803371,\n          861.5760539647758,\n          1298.3340069648193,\n          1956.497265545218,\n          2948.3026167007256,\n          4442.882938158366,\n          6695.109481084312,\n          10089.055144515105,\n          15203.490547339246,\n          22910.58196353461,\n          34524.62211052572,\n          52026.15689866611,\n          78399.72854675821,\n          118142.83050307268,\n          178032.86641424094,\n          268282.9028956321,\n          404283.3069857812,\n          609226.2702664398,\n          918060.7805699259,\n          1383452.483839961,\n          2084764.7732591254\n        ]\n      }\n    },\n    "contexts": [\n      4096,\n      131072\n    ],\n    "head_dim": 64\n  },\n  "worked_rope": {\n    "base": 10000.0,\n    "head_dim": 4,\n    "k": [\n      0.0,\n      1.0,\n      1.0,\n      -0.5\n    ],\n    "logits": [\n      1.1442484284723262,\n      1.1442484284723262,\n      1.1442484284723264\n    ],\n    "placements": [\n      [\n        3,\n        1\n      ],\n      [\n        7,\n        5\n      ],\n      [\n        2,\n        0\n      ]\n    ],\n    "q": [\n      1.0,\n      0.0,\n      0.5,\n      0.5\n    ],\n    "rotated_k_at_1": [\n      -0.8414709848078965,\n      0.5403023058681398,\n      1.0049499170837486,\n      -0.489975166874166\n    ],\n    "rotated_q_at_3": [\n      -0.9899924966004454,\n      0.1411200080598672,\n      0.48477726677324595,\n      0.5147727669757416\n    ]\n  }\n}\n')
CONFIG_BYTES = b'{"num_hidden_layers": 36, "num_experts": 128, "experts_per_token": 4, "vocab_size": 201088, "hidden_size": 2880, "intermediate_size": 2880, "swiglu_limit": 7.0, "head_dim": 64, "num_attention_heads": 64, "num_key_value_heads": 8, "sliding_window": 128, "initial_context_length": 4096, "rope_theta": 150000, "rope_scaling_factor": 32.0, "rope_ntk_alpha": 1, "rope_ntk_beta": 32}'
assert hashlib.sha256(CONFIG_BYTES).hexdigest() == '063aaa98e012cac330d58d18471700f0b0eb90bb15c58fb227809d45090f3e65'
print("Source:", 'https://huggingface.co/openai/gpt-oss-120b/resolve/b5c939de8f754692c1647ca79fbf85e8c1e70f8a/original/config.json')
config = json.loads(CONFIG_BYTES)
sc = reference['scaling']
assert (config['head_dim'], config['rope_theta'], config['initial_context_length'], config['rope_scaling_factor']) == (sc['head_dim'], sc['base'], sc['original_context'], sc['factor'])
raw = rope_inverse_frequencies(config['head_dim'], config['rope_theta'])
yarn = yarn_scaled_frequencies(config['head_dim'], config['rope_theta'], config['rope_scaling_factor'], config['initial_context_length'])
np.testing.assert_allclose(yarn.inverse_frequencies, sc['yarn']['inverse_frequencies'], rtol=1e-12, atol=1e-15)
np.testing.assert_allclose(ntk_scaled_frequencies(sc['head_dim'],sc['base'],sc['factor']),sc['ntk'],rtol=1e-12,atol=1e-15)
scaling_plot(reference, DEMO_LANGUAGE)
plt.show()
print('YaRN bounds:', yarn.low, yarn.high)
print('Q/K amplitude and full-head score factor:', yarn.attention_factor, yarn.score_factor_if_applied_to_q_and_k)

mref = reference['mrope']
x = np.array(mref['image_input'])
coords = np.array(mref['image_coordinates'], dtype=np.int64)
inv = rope_inverse_frequencies(mref['head_dim'], mref['base'])
rotated = apply_mrope_adjacent(x, coords, mref['axis_sections'], inv)
np.testing.assert_allclose(rotated,mref['image_output'],rtol=1e-12,atol=1e-12)
np.testing.assert_allclose(np.linalg.norm(rotated,axis=-1),np.linalg.norm(x,axis=-1),rtol=1e-12)
changed_coords=coords.copy();changed_coords[:,2]+=1
width_changed=apply_mrope_adjacent(x,changed_coords,mref['axis_sections'],inv)
np.testing.assert_array_equal(width_changed[:,:6],rotated[:,:6])
assert not np.allclose(width_changed[:,6:],rotated[:,6:])
text_pos=np.array([4,5,6,7],dtype=np.int64)
text_axes=np.repeat(text_pos[:,None],3,axis=1)
np.testing.assert_allclose(apply_mrope_adjacent(x,text_axes,mref['axis_sections'],inv),apply_rope_adjacent(x,text_pos,inv),rtol=1e-12,atol=1e-12)
print('Image rotations:',rotated)
print('Only width pairs change; equal text IDs recover 1D RoPE: PASS')