Module 4. Positional representations: from sinusoids to multiple axes
Read the explanations and code in the existing notebooks, or download the files. Code does not run on this site.
Download course files (.zip) 2588.1 KBThe archive preserves the source folder structure, both languages, and supporting files.
M4 practice: position, RoPE, and context extension
Main module notebook
Download .ipynb module_04_notebook_rope.ipynb · 67.5 KB
Show notebook text
M4 practice: position, RoPE, and context extension
This notebook turns the positional formulas into small CPU experiments: rotate vectors, reshape frequency spectra, and check position IDs during cached decoding.
Routes. First run all cells and interpret the results. On a second pass, change the inputs and identify which identities still hold. For reference, the code is grouped into positions, rotation, scaling, multiple axes, and caching.
Shapes. Inputs use [..., T, D], with an even width D and optional leading batch/head axes. Position IDs are integers [T]; frequencies are positive [D/2]. Numerical references use float64. The PyTorch bridge separately tests gradients and half precision, using at least float32 for phases and returning the input dtype.
Execution. No repository files, API keys, or GPU are needed. The first executable cell installs NumPy, Matplotlib, and PyTorch only if missing; installation requires a network. The remaining computation is offline. The embedded checksum detects changed reference data; mathematical correctness comes from the separate numerical checks.
0. An embedded numerical reference
These fixtures belong to M4 and do not depend on lecture wording. SHA-256 verifies the embedded JSON bytes after canonical serialization. It establishes integrity, not mathematical truth.
Python
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('Embedded Module 4 contract verified', actual_sha256)1. Frequencies and wavelengths
Pair i uses radians per position step. Its wavelength counts steps per turn. Sines occupy even coordinates and cosines occupy odd coordinates in the [T,D] table.
Inspect the slow pairs: they change most when the base increases. The diagnostic counts turns over a span L; a token sequence indexed 0…L−1 has endpoint separation L−1.
Python
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('Sinusoidal table and wavelengths verified')2. Three RoPE implementations and two pairing conventions
Adjacent pairs and HF/Llama rotate_half are not the same coordinate layout. They are linked
by a permutation that must be applied consistently to , , and projection weights.
We verify:
- adjacent matrix form = complex multiplication;
- adjacent = split-half after the explicit permutation;
- norms are preserved;
- logits differ without the permutation — the intended negative test.
Python
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('Three implementations, permutation, and norms verified', f"raw convention gap={abs(raw_adjacent-raw_split):.6f}")3. Shift invariance and position precision
With fixed vectors and frequencies, shifting both positions leaves their dot product unchanged:
We test 200 cases and reproduce the lecture's worked example. The function rope_attention_logits returns unscaled dot products. Standard attention divides by √D; at D=4, the example becomes 1.144248/2 ≈ 0.572124.
IDs remain integers until float64 phase construction. Float32 loses consecutive-integer resolution after 2²⁴; the example near 10⁸ exhibits a collision. Storing IDs as integers does not remove the precision limit of their eventual floating representation.
Python
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('Relativity and a large position ID verified', f"max_error={max_error:.3e}")4. PI, a static NTK-aware heuristic, and YaRN
- PI: divides every inverse frequency by the factor;
- static NTK-aware scaling: changes base by ;
- YaRN: blends original and interpolated spectra through a ramp.
For gpt-oss we use the continuous bounds from the original implementation (truncate=False) and
show the integer-boundary Hugging Face option (truncate=True) separately. The turn-count zones
diagnose this particular ramp; they are not a theorem about safe extrapolation. We also distinguish
the rotary attention_factor from the score factor that results when the same scaling enters both
queries and keys. This gives a global square factor only when the entire head is scaled.
Python
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/YaRN and zone geometry verified',
f"bounds={yarn.low:.4f}/{yarn.high:.4f}; "
f"zones={yarn.kept_pairs}/{yarn.transition_pairs}/{yarn.interpolated_pairs}",
)5. Read the spectrum
The upper panel compares frequencies; the lower one shows YaRN's interpolation weight. Find the unchanged first NTK pair and the fully interpolated YaRN region. These curves describe geometry, not model quality. The lecture and notebook use the same plotting function; this cell writes no files.
Python
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, 'en')
plt.show()6. ALiBi and M-RoPE
ALiBi returns bias of shape [H,T_q,T_k]; future keys receive in causal mode.
M-RoPE accepts coordinates [T,A] and assigns the D/2 frequency pairs across axes. The
image example keeps temporal ID fixed while height and width vary. Text repeats the same ID
across all three axes and is checked against ordinary one-dimensional RoPE.
Python
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 and M-RoPE geometry verified')7. Negative tests
We pin down failure modes: odd head width, wrong position length, NaN, invalid base,
incompatible Q/K, unknown pairing convention, factor below one, invalid beta ordering,
inconsistent M-RoPE sections, and non-integer coordinates.
Python
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('Negative tests passed')8. Cached positions and partial scaling
Compare full causal attention with one-token steps. New queries keep their absolute positions; cached keys are already rotated. Two negative controls reset query IDs and combine keys from one spectrum with queries from another.
Then change only an image's width coordinate and check which pair moves. Finally, test the limit of the a² rule: multiplying only the rotary part of a head need not scale the entire attention score by a².
Python
"""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())PyTorch bridge · B02 — RoPE: NumPy and PyTorch
This cell embeds the canonical NumPy→PyTorch bridge and runs its parity check. It is self-contained and does not import repository files.
Python
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")Exercises
-
Change batch size and head count while retaining
[B,H,T,D]. Explain why rotation cannot mix those leading axes. -
Reset every query position in the cache test. Locate the first mismatch and explain why the first step may still agree.
-
Build ALiBi for query positions 4, 5, 6 and key positions 0…6. Query 4 should have exactly five allowed keys.
-
Change only
win M-RoPE. Which coordinates move? With sections(2,1,1), the check above identifies coordinates 6 and 7. -
Extension. Implement Llama3-style scaling from a pinned source revision. State the wavelength boundaries and compare its frequencies with PI, NTK, and YaRN.
-
Project. Apply the experiment to an open model with pinned configuration and code. Measure quality, memory, and runtime separately. Passing the geometry checks here is not a long-context evaluation.
Module 4 project-exercise demo — Configuration-derived frequencies and multi-axis rotation
Project demonstration
Download .ipynb module_04_project_demo.ipynb · 54.3 KB
Show notebook text
Module 4 project-exercise demo — Configuration-derived frequencies and multi-axis rotation
Implementation demo, not a completed full-scale project. Exact PI/static NTK/YaRN and sectioned M-RoPE on small inputs. These are geometry checks, with no training or model-quality measurement. The gpt-oss configuration is pinned by revision and checksum. All inputs are small, deterministic, offline, and CPU-friendly. The goal is to validate interfaces, metrics, invariants, and report structure before spending GPU or API budget.
Python
from __future__ import annotations
import importlib.util, subprocess, sys
for package in ("numpy", "matplotlib"):
if importlib.util.find_spec(package) is None:
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])
import math
import random
import time
from collections import Counter, defaultdict
import matplotlib.pyplot as plt
import numpy as np
SEED = 20260801
random.seed(SEED)
np.random.seed(SEED)
print(f"reproducible smoke tier; seed={SEED}")
DEMO_LANGUAGE = 'en'
Demo 1. Recover the spectrum from a configuration
A positional spectrum can be reconstructed without loading model weights. Start by verifying the saved config.json bytes, then read the head width, base, reference length, and extension factor. The functions below build the original spectrum, PI, a static NTK-aware transformation, and the YaRN ramp. They come from the main notebook; the ALiBi and M-RoPE helpers support the next experiment.
On the plot, locate the unchanged first NTK pair and the YaRN transition. Also distinguish the amplitude applied to each rotary vector from its combined effect on a query–key dot product.
Python
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)Bounds near 8.0928 and 17.3980 yield 9 unchanged, 9 transitional, and 14 interpolated pairs. Applying an amplitude of 1.346574 to both full rotary vectors scales their dot product by about 1.813260. These results recover the positional geometry; the experiment does not measure long-text quality or establish the training history.
Demo 2. Identify which axis changed
Sectioned M-RoPE assigns its coordinate pairs to particular axes. Change only the image width coordinate and inspect which output components move. Then give a text token equal IDs on all three axes and compare it with ordinary 1D RoPE. Matching frequencies and explicit tolerances let these comparisons test specific algebraic claims.
Python
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')With sections (2,1,1), width owns the final pair, coordinates 6 and 7. Other components stay fixed, and rotations preserve vector norms. Equal text IDs explain the agreement with 1D RoPE. As a follow-up, change height instead of width and predict the affected pair before running the cell.
Scaling to the full project
For a full project, pin the model and evaluation corpus. Verify pair layout and cache positions, then measure quality across lengths, memory, and runtime separately. Positional-vector similarity is not a language-quality metric.
Minimum report
-
Data, model, and code revisions, plus the random seed.
-
The primary metric, plus quality, cost, latency, and memory where applicable.
-
A baseline, a matched budget, and at least one negative control.
-
Failed runs and the limits of the supported conclusions.