github ↗
docs menu

Module Kaun.Rope

Rotary position embeddings: a frequency schedule and its rotation.

Rotary position embeddings (Su et al., 2021).

A rotary embedding encodes a token's position by rotating pairs of features of its query and key vectors: pair i turns by the angle position * frequency.(i). The dot product of a rotated query and a rotated key then depends on their positions only through the difference, which is what attention needs.

A value of type t is the schedule: one inverse frequency per feature pair, and the cosine and sine of each pair's angle at every position below the schedule's context, computed once on the host, as every value made from numbers is; place moves it where a model computes. It has no parameters and never enters a parameter tree. Every published schedule is a choice of those frequencies, so a new one is a constructor and apply never changes.

#type 'd t

The type for rotary schedules on the set 'd: the inverse frequencies of one attention head, head_dim / 2 of them, and their angles' cosines and sines at the positions 0 to context - 1, in float32.

#val of_frequencies : context:int -> float array -> Nx.host t

of_frequencies ~context f is the schedule whose pair i has inverse frequency f.(i), for a head of 2 * Array.length f features and positions below context. Its cosines and sines are of angles computed in float64, rounded once to float32. Every other constructor is defined with it; use it for a schedule this module does not name.

Raises Invalid_argument if f is empty or holds a non-finite number, or context is not positive.

#val make : ?theta:float -> head_dim:int -> context:int -> unit -> Nx.host t

make ~head_dim ~context () is the standard schedule: pair i has frequency theta ** (-2 i / head_dim). theta defaults to 10000..

Raises Invalid_argument if head_dim is not positive and even, theta is not positive, or context is not positive.

#val llama3 : theta:float -> head_dim:int -> factor:float -> low_freq_factor:float -> high_freq_factor:float -> original_context:int -> context:int -> Nx.host t

llama3 ... is the Llama 3.1 long-context schedule: the standard frequencies with the low ones divided by factor, the high ones kept, and a linear blend between, the bands being set by the wavelengths original_context / low_freq_factor and original_context / high_freq_factor. Llama 3.1 uses ~theta:500000. ~factor:8. ~low_freq_factor:1. ~high_freq_factor:4. ~original_context:8192 ~context:131072; Llama 3.2 uses ~factor:32..

Raises Invalid_argument on a non-positive argument, or if high_freq_factor <= low_freq_factor.

#val yarn : theta:float -> head_dim:int -> factor:float -> beta_fast:float -> beta_slow:float -> original_context:int -> context:int -> Nx.host t

yarn ... is the YaRN long-context schedule (Peng et al., 2023), with the correction range left untruncated, as gpt-oss uses it. Pair i blends the standard frequency f and the interpolated one f / factor: ramp * f / factor + (1 - ramp) * f, where ramp rises linearly from 0 to 1 between the fractional pair indices that make beta_fast and beta_slow full turns over original_context positions, clamped to [0, head_dim - 1]. Fast pairs keep their frequency and slow ones are interpolated.

The schedule is its frequencies only. YaRN also prescribes an attention temperature: scores are scaled by (0.1 * ln factor + 1) ^ 2 / sqrt head_dim in place of 1 / sqrt head_dim. apply keeps norms, so a model passes that number as ?scale to Attention.scaled_dot_product_attention. Implementations that instead scale the cosines and sines by 0.1 * ln factor + 1 compute the same attention. gpt-oss uses ~theta:150000. ~factor:32. ~beta_fast:32. ~beta_slow:1. ~original_context:4096 ~context:131072.

Raises Invalid_argument if head_dim is not positive and even, theta, beta_slow, original_context or context is not positive, factor is below 1, or beta_fast <= beta_slow.

#val place : 'e Nx.Placement.t -> 'd t -> 'e t

place at t is t with its cosines and sines whole on each device of at (Nx.Placement.uncut), where a model on at reads them.

#val frequencies : 'd t -> float array

frequencies t is a copy of t's inverse frequencies, in pair order.

#val context : 'd t -> int

context t is the number of positions t rotates: positions 0 to context t - 1.

#val apply : 'd t -> pos:'d Nx.int64_t -> (float, 'b, 'd) Nx.t -> (float, 'b, 'd) Nx.t

apply t ~pos x rotates x, of shape [| batch; heads; seq; head_dim |], by the positions pos, of shape [| batch; seq |] (or [| 1; seq |], shared by the batch). Feature i is paired with feature i + head_dim / 2:

y.(i)     = x.(i) cos a     - x.(i + h) sin a
y.(i + h) = x.(i + h) cos a + x.(i) sin a

where h is head_dim / 2 and a is pos * frequency.(i).

Position 0 is the identity. A rotation keeps norms: apply never rescales x, and a schedule is a choice of angles and nothing else. The cosines and sines are the rows of t's tables at pos, cast to x's dtype for the rotation, so a float64 x is rotated to float32 accuracy. Each position must be below context t. Differentiable through Rune in x.

Raises Invalid_argument if x is not of rank 4, its last axis is not twice t's pair count, or pos has another shape.