Sinusoidal Position Embeddings kernel

[superposition]

Description

Compute sinusoidal position embeddings (fixed or learned scaling).

Operation

# Generate position embeddings for sequence
# positions: (B, T) or just T for shared positions
pos_embed = sinusoidal_embed(positions, dim=768, max_len=8192)
# Returns: (B, T, D) or (T, D)

Implementation

• Fused sin/cos computation
• Support for 2D positions (images)
• Learnable frequency scaling (Fourier features)
• RoPE-compatible output format

Tests

• Matches manual sin/cos formula
• Correct dimensionality (sin for even, cos for odd)
• 2D positions produce correct grid
• Frequencies follow geometric progression
• Works with arbitrary position values (not just 0..T)
• Output is deterministic (no randomness)

Used In

Original Transformer, NeRF, diffusion timesteps, image positional encoding