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

RotaryEmbedding

RotaryEmbedding

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

source

Bases: Module

RotaryEmbedding layer to calculate and apply the frequency tensor for complex exponentials.

Parameters:

• dim (int)
• n_heads (int)
• theta (float)
• max_seq_len (int)
• device (Device)
• head_dim (int)
• interleaved (bool)

compute_scale()

compute_scale(user_scale=None)

source

Computes the attention scale factor from head_dim or the given user scale.

Parameters:

user_scale (float | None)

Return type:

float

dim

dim: int

source

forward()

forward(x, start_pos=None, seq_len=None)

source

Applies rotary positional embeddings (RoPE) to x.

Parameters:

• x (Tensor) – Activation tensor with shape (batch, seq_len, n_kv_heads, head_dim).
• start_pos (Dim | None) – starting position of input tensor, defaults to 0 if None
• seq_len (Dim | None) – length of input tensor, defaults to x.shape[-2] if None

Returns:

Input activation tensor with rotary positional embeddings applied and the same shape as x.

Return type:

Tensor

freqs_cis

property freqs_cis: Tensor

source

Returns the flattened frequency tensor for rotary embeddings.

freqs_cis_base()

freqs_cis_base()

source

Computes the frequency tensor for complex exponentials (cis).

Tensor is scaled with theta parameter. Required to apply Rotary Position Embedding (RoPE) to tensor. See ‘Roformer: Enhanced Transformer with Rotary Embedding’ (arxiv.org/pdf/2104.09864).

Returns:

The frequency tensor for complex exponentials with shape (max_seq_len * 2, head_dim / 2, 2)

Return type:

Tensor

head_dim

head_dim: int

source

head_dim = dim // n_heads if not specified in the config.

interleaved

interleaved: bool = True

source

max_seq_len

max_seq_len: int

source

The maximum sequence length for model’s input.

n_heads

n_heads: int

source

theta

theta: float

source

Hyperparameter used to control the frequency scaling of the sinusoidal components of the embeddings.