Reference-coordinate and space-ID rotary features

Your task

Function signature

spatial_rotary_embedding(
    q: torch.Tensor,
    k: torch.Tensor,
    reference_position: torch.Tensor,
    space_uid: torch.Tensor,
    spatial_inverse_frequency: torch.Tensor,
    uid_inverse_frequency: torch.Tensor,
) -> dict[str, torch.Tensor]

Worked example

Worked example

import torch

q = torch.tensor([[[[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]]]], dtype=torch.float64, requires_grad=True)
k = torch.tensor([[[[8.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0]]]], dtype=torch.float64, requires_grad=True)
reference_position = torch.tensor([[[1.5707963267948966, 0.0, 0.0]]], dtype=torch.float64, requires_grad=True)
space_uid = torch.tensor([[2]], dtype=torch.int64)
spatial_inverse_frequency = torch.tensor([1.0], dtype=torch.float64, requires_grad=True)
uid_inverse_frequency = torch.tensor([0.7853981633974483], dtype=torch.float64, requires_grad=True)

result = spatial_rotary_embedding(q=q, k=k, reference_position=reference_position, space_uid=space_uid, spatial_inverse_frequency=spatial_inverse_frequency, uid_inverse_frequency=uid_inverse_frequency)

Expected return

{
    'q_rotated': torch.tensor([[[[-5.0, 2.0, 3.0, -8.0, 1.0000000000000002, 6.0, 7.0, 4.000000000000001]]]], dtype=torch.float64),
    'k_rotated': torch.tensor([[[[-3.9999999999999996, 7.0, 6.0, -0.9999999999999997, 8.0, 3.0, 2.0, 5.0]]]], dtype=torch.float64),
    'cosine': torch.tensor([[[6.123233995736766e-17, 1.0, 1.0, 6.123233995736766e-17, 6.123233995736766e-17, 1.0, 1.0, 6.123233995736766e-17]]], dtype=torch.float64),
    'sine': torch.tensor([[[1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0]]], dtype=torch.float64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Implement reference-position/space-UID rotary features and apply them to atom self-attention queries and keys. q and k have identical [B,H,A,D] shape with B,H,A\ge 1 and positive even D. reference_position[B,A,3] is floating xyz; space_uid[B,A] is nonnegative int64.
• spatial_inverse_frequency[S] and uid_inverse_frequency[U] are floating supplied nonnegative frequencies, with S,U\ge 0 and 3\cdot S+U\le D/2. All floating inputs share dtype (float32 or float64) and device, and all phase products are finite and representable. There is no generation of random or checkpoint frequencies.
• Construct one half-width angle vector in this exact order: x\cdot f_{s} for s=0,\ldots,S-1, then y\cdot f_{s}, then z\cdot f_{s}, then \mathrm{space\_uid}\cdot g_{u} for u=0,\ldots,U-1, then D/2-3\cdot S-U zeros. Concatenate that angle vector with itself to get [B,A,D].
• Return \mathrm{cosine}=\cos(\mathrm{angle}) and \mathrm{sine}=\sin(\mathrm{angle}), each [B,A,D].
• For each head split q into first and second halves [a,b], \mathrm{rotate\_half}(q)=[-b,a], and \mathrm{q\_rotated}=q\cdot \mathrm{cosine}[:,\mathrm{None}]+\mathrm{rotate\_half}(q)\cdot \mathrm{sine}[:,\mathrm{None}]; apply the same operation to k. Return q_rotated,k_rotated,cosine,sine.
• The inactive PAIRS span corresponding entries in both halves, not necessarily a contiguous trailing D-channel suffix. No configuration keys or masks are required. Zero-length frequency groups are allowed, including S=U=0 which leaves q/k unchanged. Noncontiguous inputs are supported; do not mutate them.
• Normalizing Q/K, choosing an atom-index attention window, and the attention operation itself are separate problems. These phases use fixed reference coordinates, not noisy sample coordinates.
• Coordinate translation with a common UID offset preserves q-k dot products because only relative phases matter; arbitrary 3D rotations need not preserve scores, so this is not a general \mathrm{SE}(3)-equivariant encoding.

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 spatial_rotary_embedding(
    q: torch.Tensor,  # (B, H, A, D)
    k: torch.Tensor,  # (B, H, A, D)
    reference_position: torch.Tensor,  # (B, A, 3)
    space_uid: torch.Tensor,  # (B, A)
    spatial_inverse_frequency: torch.Tensor,  # (S,)
    uid_inverse_frequency: torch.Tensor,  # (U,)
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    q_rotated: (B, H, A, D).
    k_rotated: (B, H, A, D).
    cosine: (B, A, D).
    sine: (B, A, D)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement spatial_rotary_embedding")

Sources and further reading

• www.biorxiv.org
• huggingface/transformers · modeling_esmfold2.py