Rotary positions for queries and keys

Your task

Function signature

rotary_position_embedding(
    q: torch.Tensor,
    k: torch.Tensor,
    q_positions: torch.Tensor,
    k_positions: torch.Tensor,
    inverse_frequency: torch.Tensor,
    options: dict,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Quarter-turn, half-turn, and unchanged suffix

import torch

q = torch.tensor([[[[1.0, 2.0, 3.0, 4.0, 9.0]]]], dtype=torch.float64, requires_grad=True)
k = torch.tensor([[[[2.0, 1.0, 4.0, 3.0, 8.0]]]], dtype=torch.float64, requires_grad=True)
q_positions = torch.tensor([[1]], dtype=torch.int64)
k_positions = torch.tensor([[0]], dtype=torch.int64)
inverse_frequency = torch.tensor([1.5707963267948966, 3.141592653589793], dtype=torch.float64, requires_grad=True)
options = {
    'rotary_dim': 4,
    'convention': 'adjacent',
}

result = rotary_position_embedding(q=q, k=k, q_positions=q_positions, k_positions=k_positions, inverse_frequency=inverse_frequency, options=options)

Expected return

(torch.tensor([[[[-2.0, 1.0000000000000002, -3.0000000000000004, -3.9999999999999996, 9.0]]]], dtype=torch.float64), torch.tensor([[[[2.0, 1.0, 4.0, 3.0, 8.0]]]], dtype=torch.float64))

Parameters

Return

tuple[torch.Tensor, torch.Tensor]

Required behavior

• Rotate the first R features of q[B,H,L_q,D] and k[B,H,L_k,D] without changing the suffix. B,H,L_q,L_k\,\ge \,1. Positions are supplied separately as integer q_positions[B,L_q] and k_positions[B,L_k]; different lengths and offsets model cached keys.
• inverse_frequency[R/2] is supplied, shared by both tensors, and has their floating dtype/device. All tensor values are finite. Config: rotary_dim R defaults to D and must be even with 2\,\le \,R\,\le \,D; convention is 'half_split' (default) or 'adjacent'. There are no other options.
• For pair r, \theta[b,t,r]\,=\,\mathrm{positions}[b,t]\,\cdot \,\mathrm{inverse\_frequency}[r]. Half-split pairs coordinates (r,\,r+R/2); adjacent pairs (2r,\,2r+1).
• In either convention, (a,b)\,\mapsto \,(a\cdot \cos(\theta)-b\cdot \sin(\theta),\,a\cdot \sin(\theta)+b\cdot \cos(\theta)). Do not mix the two pair layouts or reverse the sine sign. Angles broadcast over H, not across positions or batches.
• The suffix x[..., R:] is unchanged.

Why this matters

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def rotary_position_embedding(
    q: torch.Tensor,  # (B, H, Lq, D)
    k: torch.Tensor,  # (B, H, Lk, D)
    q_positions: torch.Tensor,  # (B, Lq)
    k_positions: torch.Tensor,  # (B, Lk)
    inverse_frequency: torch.Tensor,  # (R/2)
    options: dict,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    q_rotated: (B, H, Lq, D).
    k_rotated: (B, H, Lk, D)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement rotary_position_embedding")

Sources and further reading

• facebookresearch/esm · rotary_embedding.py
• Biohub/esm · layers.py
• salesforce/progen · modeling_progen.py