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

TransposedRotaryEmbedding

TransposedRotaryEmbedding

class max.experimental.nn.rope.TransposedRotaryEmbedding(weight)

source

Bases: RotaryEmbedding

Applies RoPE using a transposed head-dimension layout.

Parameters:

weight (Tensor)

forward()

forward(x, start_pos=0)

source

Applies rotary positional embeddings (RoPE) to x.

The representation of x is transposed within the final dimension compared to traditional RotaryEmbedding.

seq_len is inferred from the shape of x.

Parameters:

• x (Tensor) – Activation tensor with shape (batch, seq_len, n_kv_heads, head_dim). x is interpreted as a complex number valued tensor where the first half of head_dim are the real parts and the last half are the imaginary parts.
• start_pos (int | str | Dim | integer | TypedAttr) – starting position of input tensor, defaults to 0 if None

Returns:

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

Return type:

Tensor