RotaryEmbedding
ModuleRotaryEmbedding(head_dim: int, max_position_embeddings: int, base: float = 10000.0)Precomputed cos / sin tables for rotary positional embedding.
Owns no learnable parameters — a thin wrapper around two registered
buffers so the tables move with .to(device=...) and serialise with
the rest of the model state. forward() returns the precomputed
(cos, sin) pair; callers pass them into
lucid.nn.functional.apply_rotary_emb along with the query /
key tensors.
Parameters
head_dimintmax_position_embeddingsintbasefloat= 10000.010000.0 per the RoFormer / LLaMA / GPT-NeoX
convention; some long-context models (e.g. CodeLlama) use
1_000_000.Notes
The cos_cached / sin_cached tables are built once at
construction and registered as non-persistent buffers. They follow
.to(device=...) automatically but are not saved in state_dict
(RoPE has no learnable state — regenerate at load time). The module
form is the right choice for any transformer that reuses the same
max_position_embeddings across calls; the functional
lucid.nn.functional.apply_rotary_emb consumes the cached
pair and applies the rotation in place to q and k. See
that function for the rotation math.
Examples
>>> import lucid.nn as nn
>>> rope = nn.RotaryEmbedding(head_dim=64, max_position_embeddings=2048)
>>> cos, sin = rope()
>>> cos.shape, sin.shape
((2048, 64), (2048, 64))