For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

YarnRotaryEmbedding

YarnRotaryEmbedding

class max.experimental.nn.common_layers.rotary_embedding.YarnRotaryEmbedding(dim, n_heads, theta, max_seq_len, device, head_dim=None, interleaved=True, scaling_params=None)

source

Bases: RotaryEmbedding

Generic YaRN (Yet another RoPE eNhancement) Rotary Position Embedding layer.

This implementation provides YARN scaling for models that require it, with configurable parameters for beta_fast, beta_slow, and scaling factor.

Initialize YarnRotaryEmbedding.

Parameters:

• dim (int) – The dimension of the rotary embedding (usually hidden_size).
• n_heads (int) – Number of attention heads.
• theta (float) – Base frequency for rotary embeddings.
• max_seq_len (int) – Maximum sequence length.
• device (Device) – Device to place the embeddings on.
• head_dim (int) – Optional head dimension override.
• interleaved (bool) – Whether to use interleaved complex format.
• scaling_params (YarnScalingParams | None) – YARN scaling parameters.

freqs_cis_base()

freqs_cis_base()

source

Computes the frequency tensor for complex exponentials (cis) with YARN scaling applied.

Return type:

Tensor

scaling_params

scaling_params: YarnScalingParams | None = None

source