Registry / ai-ml / rotary-embedding-torch

rotary-embedding-torch

library0.9.1pypypi✓ verified 24d ago

This library provides a Pytorch implementation of the Rotary Positional Embedding (RoPE), a crucial component for modern transformer architectures like LLaMA, designed to improve the model's ability to handle long sequences. It offers an easy-to-use API to apply rotary embeddings to query and key tensors. The current version is 0.8.9, and it follows a rapid release cadence for bug fixes and minor improvements.

ai-ml

pip install rotary-embedding-torch

github ↗package ↗

INSTALL

IMPORT

SIG · ROTARY-EMBEDDING-T

rotary-embedding-torch

ai-mlpythonv0.9.1

Install

75.4s avg

Import

5436ms

Disk

5079MB

Pass rate

5/ 10

Env Coverage5 / 10

glibc

3.9–3.13

musl

3.9–3.13

Install & Compatibility

Where this runs

tested against v0.9.1 · pip install

no network on importno background threads

Install × environment matrix

Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux

musl

py 3.10–3.95 runs

build_error

glibc

py 3.10–3.95 runs

installs and imports cleanly · install 75.4s · import 5.436s · 4710.4MB

5079MB installed

● package 5079MB

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

RotaryEmbedding

✓ from rotary_embedding_torch import RotaryEmbedding

This example demonstrates how to initialize `RotaryEmbedding` and apply it to example query and key tensors, typically used in self-attention mechanisms. The `dim` parameter must match the last dimension of your input tensors (head_dim).

import torch from rotary_embedding_torch import RotaryEmbedding # Define embedding dimension dim = 64 # Initialize RotaryEmbedding. max_seq_len can be set for pre-computation. # If not set, it's computed dynamically based on input. rotary_emb = RotaryEmbedding(dim=dim, max_seq_len=2048) # Create dummy query and key tensors # shape: (batch_size, num_heads, sequence_length, head_dim) seq_len = 1024 q = torch.randn(1, 8, seq_len, dim) k = torch.randn(1, 8, seq_len, dim) # Apply rotary embeddings q_rot = rotary_emb(q) k_rot = rotary_emb(k) print(f"Original query shape: {q.shape}") print(f"Rotary-embedded query shape: {q_rot.shape}") print(f"Example embedded value (first element): {q_rot[0, 0, 0, 0].item():.4f}")

Debug

Known issues

gotchaPrior to v0.8.0, there was a bug in the chi scale multiplication which could lead to incorrect positional embeddings, particularly when using specific scaling factors. Models trained with affected versions might show subtle performance degradation or instability.

fix

Upgrade to `rotary-embedding-torch>=0.8.0`.

affects: < 0.8.0

gotchaWhen using `torch.compile` for performance optimization, versions prior to v0.8.6 might encounter issues due to the `seq_len` being cached as a non-integer type. This could lead to compilation failures or incorrect behavior with JIT.

fix

Upgrade to `rotary-embedding-torch>=0.8.6` to ensure compatibility with `torch.compile`.

affects: < 0.8.6

gotchaThe `RotaryEmbedding` object caches internal calculations based on the maximum sequence length encountered. If inputs change drastically in `seq_len` (e.g., during inference with varying sequence lengths) or if `max_seq_len` is not adequately pre-configured, it might lead to unnecessary re-computations or out-of-bounds errors if the input `seq_len` exceeds the initially cached maximum.

fix

Initialize `RotaryEmbedding` with an appropriate `max_seq_len` if you know the maximum sequence length your model will handle to optimize caching and prevent dynamic re-computation overhead.

affects: All versions

Errors

Common errors & fixes

ModuleNotFoundError: No module named 'rotary_embedding_torch'

The 'rotary-embedding-torch' package is not installed in the current Python environment.

fix

pip install rotary-embedding-torch

TypeError: RotaryEmbedding.__init__() missing 1 required positional argument: 'dim'

The 'RotaryEmbedding' class constructor requires the 'dim' argument, specifying the feature dimension of the tensors it will process.

fix

from rotary_embedding_torch import RotaryEmbedding

rope = RotaryEmbedding(dim = 128) # Replace 128 with your actual feature dimension

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

The 'RotaryEmbedding' instance and the input tensors (query and key) must be explicitly moved to the same PyTorch device (e.g., all on CPU or all on CUDA).

fix

from rotary_embedding_torch import RotaryEmbedding
import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

rope = RotaryEmbedding(dim = 128).to(device)
q = torch.randn(1, 8, 32, 128).to(device)
k = torch.randn(1, 8, 32, 128).to(device)

q, k = rope(q, k)

AttributeError: 'RotaryEmbedding' object has no attribute 'apply_rotary_pos_emb'

The 'RotaryEmbedding' instance is a callable object (an 'nn.Module') and rotary embeddings are applied by calling the instance directly with query and key tensors, not by calling an explicit method like 'apply_rotary_pos_emb' on the instance.

fix

from rotary_embedding_torch import RotaryEmbedding
import torch

rope = RotaryEmbedding(dim = 128)
q = torch.randn(1, 8, 32, 128)
k = torch.randn(1, 8, 32, 128)

# Correct way to apply rotary embeddings
q, k = rope(q, k)

Upgrade

Version history

0.9.1latest on PyPI · released Jun 20, 2026

Audit

Dependencies

torchrequiredCore deep learning framework

Agent activity

16 hits · last 30 days

node

OpenAI (training)

Resources

githubhttps://github.com/lucidrains/rotary-embedding-torch ↗

packagehttps://pypi.org/project/rotary-embedding-torch/ ↗