RotaryEmbedding#

class sdm.nn.RotaryEmbedding(channels: int, layout: Literal['split_half', 'interleaved'], theta: float = 100000, requires_grad: bool = True, partial_rotary_factor: float = 1.0, device: device | str | None = None, dtype: dtype | None = None)#

Bases: Module

Rotary Positional Embedding (RoPE) from the “RoFormer: Enhanced Transformer with Rotary Position Embedding” paper.

Parameters:

  • channels (int) – The number of channels per attention head.

  • layout (Literal['split_half', 'interleaved']) – The channel pairing layout. "split_half" pairs the first half of the channels with the second half. "interleaved" pairs adjacent even and odd channels.

  • theta (float) – The base frequency used to initialize inverse frequencies.

  • requires_grad (bool) – Whether inverse frequencies are learnable.

  • partial_rotary_factor (float) – The fraction of leading channels to which RoPE is applied.

  • device (device | str | None) – The device.

  • dtype (dtype | None) – The dtype.

forward(x: Tensor) → Tensor#

The forward pass.

Parameters:

x (Tensor) – Tensor with shape [..., S, H, C]. S is the query sequence length, H is the number of attention heads, and C is the channels per head.

Returns:

Tensor with shape [..., S, H, C].

Return type:

Tensor