rotary

rotary

RotaryPositionalEncoding

RotaryPositionalEncoding(embedding_dimension, number_of_heads, base_frequency=10000.0, learnable_frequencies=False)

Bases: Module

Base class for rotary positional encoding, from https://arxiv.org/pdf/2104.09864.

Initialize rotary positional encoding frequencies.

Parameters:

Name	Type	Description	Default
embedding_dimension	int	Full model embedding dimension.	required
number_of_heads	int	Number of attention heads.	required
base_frequency	float	Base frequency for geometric spacing.	10000.0
learnable_frequencies	bool	Whether frequency bands are trainable.	False

Raises:

Type	Description
ValueError	If attention dimensions or frequencies are invalid.

Source code in src/versatil/models/layers/positional_encoding/rotary.py

def __init__(
    self,
    embedding_dimension: int,
    number_of_heads: int,
    base_frequency: float = 10000.0,
    learnable_frequencies: bool = False,
):
    """Initialize rotary positional encoding frequencies.

    Args:
        embedding_dimension: Full model embedding dimension.
        number_of_heads: Number of attention heads.
        base_frequency: Base frequency for geometric spacing.
        learnable_frequencies: Whether frequency bands are trainable.

    Raises:
        ValueError: If attention dimensions or frequencies are invalid.
    """
    super().__init__()
    if embedding_dimension <= 0:
        raise ValueError(
            f"embedding_dimension must be positive, got {embedding_dimension}."
        )
    if number_of_heads <= 0:
        raise ValueError(
            f"number_of_heads must be positive, got {number_of_heads}."
        )
    if embedding_dimension % number_of_heads != 0:
        raise ValueError(
            f"embedding_dimension ({embedding_dimension}) must be divisible "
            f"by number_of_heads ({number_of_heads})."
        )
    if base_frequency <= 0.0:
        raise ValueError(f"base_frequency must be positive, got {base_frequency}.")
    self.embedding_dimension = embedding_dimension
    self.number_of_heads = number_of_heads
    self.head_dimension = embedding_dimension // number_of_heads

    if self.head_dimension % 2 != 0:
        raise ValueError("head_dimension must be even for rotary encoding")

    frequencies = self._compute_frequencies(
        dimension=self.head_dimension, base_frequency=base_frequency
    )
    self.register_parameter(
        "frequencies",
        nn.Parameter(frequencies, requires_grad=learnable_frequencies),
    )

apply_rotation staticmethod

apply_rotation(tensor, sine, cosine)

Apply rotary transformation using interleaved even/odd convention.

Pairs even and odd indices: [-odd, even] interleaved back. This is the original RoFormer/VersatIL convention.

Parameters:

Name	Type	Description	Default
tensor	Tensor	Input (B, number_of_heads, L, head_dim).	required
sine	Tensor	Sine components matching sequence + head_dim shape.	required
cosine	Tensor	Cosine components matching sequence + head_dim shape.	required

Returns:

Type	Description
Tensor	Rotated tensor of same shape.

Source code in src/versatil/models/layers/positional_encoding/rotary.py

@staticmethod
def apply_rotation(
    tensor: torch.Tensor, sine: torch.Tensor, cosine: torch.Tensor
) -> torch.Tensor:
    """Apply rotary transformation using interleaved even/odd convention.

    Pairs even and odd indices: ``[-odd, even]`` interleaved back.
    This is the original RoFormer/VersatIL convention.

    Args:
        tensor: Input (B, number_of_heads, L, head_dim).
        sine: Sine components matching sequence + head_dim shape.
        cosine: Cosine components matching sequence + head_dim shape.

    Returns:
        Rotated tensor of same shape.
    """
    even_indices = tensor[..., 0::2]
    odd_indices = tensor[..., 1::2]
    rotated_pairs = torch.stack([-odd_indices, even_indices], dim=-1)
    rotated_pairs = rotated_pairs.flatten(-2)
    rotated_tensor = (tensor * cosine) + (rotated_pairs * sine)
    return rotated_tensor

apply_rotation_half staticmethod

apply_rotation_half(tensor, sine, cosine)

Apply rotary transformation using split-half convention.

Splits the last dimension in half: [-second_half, first_half].

Parameters:

Name	Type	Description	Default
tensor	Tensor	Input (B, number_of_heads, L, head_dim).	required
sine	Tensor	Sine components matching sequence + head_dim shape.	required
cosine	Tensor	Cosine components matching sequence + head_dim shape.	required

Returns:

Type	Description
Tensor	Rotated tensor of same shape.

Source code in src/versatil/models/layers/positional_encoding/rotary.py

@staticmethod
def apply_rotation_half(
    tensor: torch.Tensor, sine: torch.Tensor, cosine: torch.Tensor
) -> torch.Tensor:
    """Apply rotary transformation using split-half convention.

    Splits the last dimension in half: ``[-second_half, first_half]``.

    Args:
        tensor: Input (B, number_of_heads, L, head_dim).
        sine: Sine components matching sequence + head_dim shape.
        cosine: Cosine components matching sequence + head_dim shape.

    Returns:
        Rotated tensor of same shape.
    """
    first_half = tensor[..., : tensor.shape[-1] // 2]
    second_half = tensor[..., tensor.shape[-1] // 2 :]
    rotated = torch.cat([-second_half, first_half], dim=-1)
    return (tensor * cosine) + (rotated * sine)

RotaryPositionalEncoding1D

RotaryPositionalEncoding1D(embedding_dimension, number_of_heads, base_frequency=10000.0, learnable_frequencies=False)

Bases: RotaryPositionalEncoding

Rotary positional encoding for 1D sequences.

Source code in src/versatil/models/layers/positional_encoding/rotary.py

def __init__(
    self,
    embedding_dimension: int,
    number_of_heads: int,
    base_frequency: float = 10000.0,
    learnable_frequencies: bool = False,
):
    """Initialize rotary positional encoding frequencies.

    Args:
        embedding_dimension: Full model embedding dimension.
        number_of_heads: Number of attention heads.
        base_frequency: Base frequency for geometric spacing.
        learnable_frequencies: Whether frequency bands are trainable.

    Raises:
        ValueError: If attention dimensions or frequencies are invalid.
    """
    super().__init__()
    if embedding_dimension <= 0:
        raise ValueError(
            f"embedding_dimension must be positive, got {embedding_dimension}."
        )
    if number_of_heads <= 0:
        raise ValueError(
            f"number_of_heads must be positive, got {number_of_heads}."
        )
    if embedding_dimension % number_of_heads != 0:
        raise ValueError(
            f"embedding_dimension ({embedding_dimension}) must be divisible "
            f"by number_of_heads ({number_of_heads})."
        )
    if base_frequency <= 0.0:
        raise ValueError(f"base_frequency must be positive, got {base_frequency}.")
    self.embedding_dimension = embedding_dimension
    self.number_of_heads = number_of_heads
    self.head_dimension = embedding_dimension // number_of_heads

    if self.head_dimension % 2 != 0:
        raise ValueError("head_dimension must be even for rotary encoding")

    frequencies = self._compute_frequencies(
        dimension=self.head_dimension, base_frequency=base_frequency
    )
    self.register_parameter(
        "frequencies",
        nn.Parameter(frequencies, requires_grad=learnable_frequencies),
    )

compute_rotation_components

compute_rotation_components(seq_len)

Computes sine and cosine components for 1D sequence positions.

Parameters:

Name	Type	Description	Default
seq_len	int	Sequence length.	required

Returns:

Type	Description
tuple[Tensor, Tensor]	Tuple of (sine, cosine) tensors of shape (seq_len, head_dim).

Source code in src/versatil/models/layers/positional_encoding/rotary.py

def compute_rotation_components(
    self, seq_len: int
) -> tuple[torch.Tensor, torch.Tensor]:
    """Computes sine and cosine components for 1D sequence positions.

    Args:
        seq_len: Sequence length.

    Returns:
        Tuple of (sine, cosine) tensors of shape (seq_len, head_dim).
    """
    device = self.frequencies.device
    position_indices = torch.arange(seq_len, device=device)
    angles = position_indices[:, None] * self.frequencies[None, :]
    sine_components = torch.sin(angles)
    cosine_components = torch.cos(angles)
    return sine_components, cosine_components

RasterRotaryPositionalEncoding2D

RasterRotaryPositionalEncoding2D(embedding_dimension, number_of_heads, base_frequency=10000.0, learnable_frequencies=False)

Bases: RotaryPositionalEncoding

Rotary encoding over flattened raster positions of a 2D grid.

Matches the DFormerv2 reference convention: every token is rotated by its flattened index row * width + column with a single frequency band spanning the full head dimension, spaced as 1 / base_frequency ** linspace(0, 1, head_dim // 2) with the endpoint included. Pretrained DFormerv2 checkpoints require exactly this scheme.

Initialize raster rotary encoding with endpoint-spaced frequencies.

Parameters:

Name	Type	Description	Default
embedding_dimension	int	Full model embedding dimension.	required
number_of_heads	int	Number of attention heads.	required
base_frequency	float	Base frequency for geometric spacing.	10000.0
learnable_frequencies	bool	Whether frequency bands are trainable.	False

Source code in src/versatil/models/layers/positional_encoding/rotary.py

def __init__(
    self,
    embedding_dimension: int,
    number_of_heads: int,
    base_frequency: float = 10000.0,
    learnable_frequencies: bool = False,
):
    """Initialize raster rotary encoding with endpoint-spaced frequencies.

    Args:
        embedding_dimension: Full model embedding dimension.
        number_of_heads: Number of attention heads.
        base_frequency: Base frequency for geometric spacing.
        learnable_frequencies: Whether frequency bands are trainable.
    """
    super().__init__(
        embedding_dimension=embedding_dimension,
        number_of_heads=number_of_heads,
        base_frequency=base_frequency,
        learnable_frequencies=learnable_frequencies,
    )
    half_dimension = self.head_dimension // 2
    exponents = torch.linspace(0, 1, half_dimension)
    frequencies = 1.0 / (base_frequency**exponents)
    with torch.no_grad():
        self.frequencies.copy_(frequencies.repeat_interleave(2))

compute_rotation_components

compute_rotation_components(height, width)

Computes sine and cosine components for raster grid positions.

Parameters:

Name	Type	Description	Default
height	int	Grid height.	required
width	int	Grid width.	required

Returns:

Type	Description
tuple[Tensor, Tensor]	Tuple of (sine, cosine) tensors of shape (H, W, head_dim).

Source code in src/versatil/models/layers/positional_encoding/rotary.py

def compute_rotation_components(
    self, height: int, width: int
) -> tuple[torch.Tensor, torch.Tensor]:
    """Computes sine and cosine components for raster grid positions.

    Args:
        height: Grid height.
        width: Grid width.

    Returns:
        Tuple of (sine, cosine) tensors of shape (H, W, head_dim).
    """
    device = self.frequencies.device
    position_indices = torch.arange(height * width, device=device)
    angles = position_indices[:, None] * self.frequencies[None, :]
    sine_components = torch.sin(angles).reshape(height, width, -1)
    cosine_components = torch.cos(angles).reshape(height, width, -1)
    return sine_components, cosine_components

RotaryPositionalEncoding2D

RotaryPositionalEncoding2D(embedding_dimension, number_of_heads, base_frequency=10000.0, learnable_frequencies=False)

Bases: RotaryPositionalEncoding

Rotary positional encoding for 2D spatial grids.

Initialize rotary positional encoding for 2D grids.

Parameters:

Name	Type	Description	Default
embedding_dimension	int	Full model embedding dimension.	required
number_of_heads	int	Number of attention heads.	required
base_frequency	float	Base frequency for geometric spacing.	10000.0
learnable_frequencies	bool	Whether frequency bands are trainable.	False

Raises:

Type	Description
ValueError	If per-axis head dimensions are invalid.

Source code in src/versatil/models/layers/positional_encoding/rotary.py

def __init__(
    self,
    embedding_dimension: int,
    number_of_heads: int,
    base_frequency: float = 10000.0,
    learnable_frequencies: bool = False,
):
    """Initialize rotary positional encoding for 2D grids.

    Args:
        embedding_dimension: Full model embedding dimension.
        number_of_heads: Number of attention heads.
        base_frequency: Base frequency for geometric spacing.
        learnable_frequencies: Whether frequency bands are trainable.

    Raises:
        ValueError: If per-axis head dimensions are invalid.
    """
    super().__init__(
        embedding_dimension=embedding_dimension,
        number_of_heads=number_of_heads,
        base_frequency=base_frequency,
        learnable_frequencies=learnable_frequencies,
    )
    self.half_head_dim = self.head_dimension // 2
    if self.half_head_dim % 2 != 0:
        raise ValueError("half_head_dimension must be even for 2D rotary encoding")
    freq_set = self._compute_frequencies(
        self.half_head_dim, base_frequency=base_frequency
    )
    with torch.no_grad():
        self.frequencies.copy_(torch.cat([freq_set, freq_set]))

compute_rotation_components

compute_rotation_components(height, width)

Computes sine and cosine components for 2D grid positions.

Parameters:

Name	Type	Description	Default
height	int	Grid height.	required
width	int	Grid width.	required

Returns:

Type	Description
tuple[Tensor, Tensor]	Tuple of (sine, cosine) tensors of shape (H, W, head_dim).

Source code in src/versatil/models/layers/positional_encoding/rotary.py

def compute_rotation_components(
    self, height: int, width: int
) -> tuple[torch.Tensor, torch.Tensor]:
    """Computes sine and cosine components for 2D grid positions.

    Args:
        height: Grid height.
        width: Grid width.

    Returns:
        Tuple of (sine, cosine) tensors of shape (H, W, head_dim).
    """
    # Use frequencies' device - parameters are automatically moved when module.to(device) is called
    device = self.frequencies.device
    # y positions (rows), x positions (cols)
    y_pos = torch.arange(height, device=device)[:, None].repeat(1, width)
    x_pos = torch.arange(width, device=device)[None, :].repeat(height, 1)
    # Split frequencies: first half for y, second for x
    freq_y = self.frequencies[: self.half_head_dim]
    freq_x = self.frequencies[self.half_head_dim :]
    angles_y = y_pos[..., None] * freq_y[None, None, :]
    angles_x = x_pos[..., None] * freq_x[None, None, :]
    angles = torch.cat([angles_y, angles_x], dim=-1)
    sine_components = torch.sin(angles)
    cosine_components = torch.cos(angles)
    return sine_components, cosine_components