### Configure Device Placement

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/configuration.md

Examples for placing the module on specific GPUs or using DataParallel for multi-GPU setups.

```python
# On specific GPU
device = 'cuda:0'
q = torch.randn(..., device=device)
rotary_emb = RotaryEmbedding(dim=64).to(device)

# Multi-GPU with DataParallel
model = MyModel().to(device)
model = nn.DataParallel(model)
```

--------------------------------

### Tuple Unpacking Example

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/types.md

Demonstrates how multi-output methods return tuples that are unpacked into individual variables.

```python
q_rot, k_rot = rotary_emb.rotate_queries_and_keys(q, k)
```

--------------------------------

### Configure Precision for Training

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/configuration.md

Examples for initializing tensors and modules with different precision types.

```python
# float32 (default, recommended for stability)
rotary_emb = RotaryEmbedding(dim=64)
q = torch.randn(..., dtype=torch.float32)

# float16 (less stable, requires careful scaling)
rotary_emb = RotaryEmbedding(dim=64)
q = torch.randn(..., dtype=torch.float16)

# bfloat16 (good stability, limited hardware support)
rotary_emb = RotaryEmbedding(dim=64)
q = torch.randn(..., dtype=torch.bfloat16)
```

--------------------------------

### Install rotary-embedding-torch

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/README.md

Install the library via pip.

```bash
$ pip install rotary-embedding-torch
```

--------------------------------

### Static Type Checking with RotaryEmbedding

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/types.md

Example demonstrating the use of type annotations for integration with static type checkers like mypy, pyright, and pylance.

```python
# Type-checking with mypy, pyright, pylance:
import torch
from rotary_embedding_torch import RotaryEmbedding

def my_attention(q: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
    rotary_emb: RotaryEmbedding = RotaryEmbedding(dim=64)
    q_rot: torch.Tensor = rotary_emb.rotate_queries_or_keys(q)
    k_rot: torch.Tensor = rotary_emb.rotate_queries_or_keys(k)
    return q_rot
```

--------------------------------

### Function Signature Type Annotations

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/types.md

Example of type hinting for the RotaryEmbedding constructor parameters.

```python
def __init__(
    self,
    dim: int,
    custom_freqs: Tensor | None = None,
    freqs_for: Literal['lang', 'pixel', 'constant'] = 'lang',
    ...
) -> None
```

--------------------------------

### Implement Cache Fetch and Store

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/helpers-and-internals.md

Retrieves cached frequencies if within range or computes and stores new frequencies when starting from offset zero.

```python
# Check if cached
if should_cache and exists(self.cached_freqs) and (seq_len + offset) <= self.cached_freqs_seq_len:
    return self.cached_freqs[offset:(offset + seq_len)].detach()

# Compute if not cached
freqs = compute_frequencies()

# Store in cache
if should_cache and offset == 0:
    self.cached_freqs[:seq_len] = freqs.detach()
    self.cached_freqs_seq_len = seq_len
```

--------------------------------

### Python 3.10+ Union Type Syntax

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/types.md

Examples of using the pipe operator for type annotations, equivalent to older Optional or Union types.

```python
Tensor | None                              # Either Tensor or None
int | float                                # Either int or float
tuple[int | float, ...] | Tensor | None   # Multiple union possibilities
```

--------------------------------

### Basic Multi-Head Attention Rotation

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/apply-rotary-emb.md

Demonstrates applying rotary embeddings to query and key tensors in a standard multi-head attention setup.

```python
import torch
from rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb

# Setup
rotary_emb = RotaryEmbedding(dim=64)
seq_len = 512
batch_size = 2
num_heads = 8

# Compute frequencies once
freqs = rotary_emb(torch.arange(seq_len))

# Create attention tensors
q = torch.randn(batch_size, num_heads, seq_len, 64)
k = torch.randn(batch_size, num_heads, seq_len, 64)

# Apply rotations
q_rot = apply_rotary_emb(freqs, q)
k_rot = apply_rotary_emb(freqs, k)

# Continue with attention...
sim = torch.einsum('b h i d, b h j d -> b h i j', q_rot, k_rot)
```

--------------------------------

### Initialize XPos Scale Coefficients

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/helpers-and-internals.md

Calculates initial scale coefficients per frequency pair using a normalized linear shift.

```python
scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)
```

--------------------------------

### Run Type Checking Tools

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/types.md

Command-line tools for validating type hints in your project.

```bash
mypy your_code.py
pyright your_code.py
pylance (in VS Code)
```

--------------------------------

### Compare Broadcat with Standard PyTorch

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/broadcat.md

Demonstrates the convenience of broadcat compared to manually calling broadcast_tensors followed by cat.

```python
# Traditional approach
broadcasted = torch.broadcast_tensors(*tensors)
result = torch.cat(broadcasted, dim=dim)

# Using broadcat
result = broadcat(tensors, dim=dim)
```

--------------------------------

### Rotate Queries and Keys with XPOS

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/index.md

Use the correct method when XPOS is enabled to avoid assertion errors.

```python
# Wrong:
q_rot = rotary_emb.rotate_queries_or_keys(q)  # AssertionError

# Right:
q_rot, k_rot = rotary_emb.rotate_queries_and_keys(q, k)
```

--------------------------------

### Project Directory Structure

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/index.md

Overview of the repository file organization.

```text
rotary-embedding-torch/
├── rotary_embedding_torch/
│   ├── __init__.py                 # Public API exports
│   ├── rotary_embedding_torch.py   # Core implementation
│   └── flash_attn_with_rotary.py   # Fused attention kernel
├── tests/
│   └── test_rotary.py              # Test suite
├── README.md                        # User documentation
├── pyproject.toml                   # Package metadata
└── train_cifar100.py                # Example usage
```

--------------------------------

### Component Documentation Template

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/README.md

Standard structure for documenting individual components, including overview, signatures, and usage.

```markdown
# Component Name

## Overview
- What it does
- Where to import it
- Key features

## Function/Method Signature
- Full Python signature
- Parameter table with types and descriptions
- Return type documentation

## Usage Examples
- Basic usage
- Common patterns
- Edge cases

## Related Components
- Links to related documentation
- Dependencies or interactions

## Source Reference
- File location
- Line numbers
```

--------------------------------

### apply_rotary_emb(freqs, t, start_index=0, scale=1, seq_dim=-2)

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/apply-rotary-emb.md

Applies rotary positional embeddings to a given tensor. The function rotates a subset of features defined by the frequency tensor and optional start index, while leaving other features unchanged.

```APIDOC
## apply_rotary_emb(freqs, t, start_index=0, scale=1, seq_dim=-2)

### Description
Applies rotary positional embeddings to the input tensor `t` using the provided `freqs`. Only the features starting at `start_index` up to the length of the frequency dimension are rotated, while others remain unchanged.

### Parameters
- **freqs** (torch.Tensor) - Required - Pre-computed frequency tensor.
- **t** (torch.Tensor) - Required - The input tensor to be rotated.
- **start_index** (int) - Optional - The starting feature index for rotation. Defaults to 0.
- **scale** (float) - Optional - Scaling factor applied to the rotation. Defaults to 1.
- **seq_dim** (int) - Optional - The dimension corresponding to the sequence length. Defaults to -2.

### Usage Example
```python
from rotary_embedding_torch import apply_rotary_emb

# Apply rotation to tensor q
q_rot = apply_rotary_emb(freqs, q, start_index=0, scale=0.95)
```
```

--------------------------------

### Configure Length Extrapolation (XPos)

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/configuration.md

Enables position-dependent decay scaling for length extrapolation. Requires using rotate_queries_and_keys() instead of rotate_queries_or_keys().

```python
# Standard rotary embeddings
rotary_emb = RotaryEmbedding(dim=64, use_xpos=False)
q_rot = rotary_emb.rotate_queries_or_keys(q)

# Length-extrapolatable embeddings
rotary_emb = RotaryEmbedding(dim=64, use_xpos=True)
q_rot, k_rot = rotary_emb.rotate_queries_and_keys(q, k)
```

--------------------------------

### Import broadcat

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/broadcat.md

Import the broadcat function from the rotary_embedding_torch package.

```python
from rotary_embedding_torch import broadcat
```

--------------------------------

### Apply Basic Learned Rotations

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/apply-learned-rotations.md

Demonstrates the fundamental usage of applying learnable rotation parameters to a tensor.

```python
import torch
from rotary_embedding_torch import apply_learned_rotations

# Create learnable rotation parameters
num_rotation_params = 32
rotations = torch.nn.Parameter(
    torch.randn(num_rotation_params) * 0.1
)

# Tensor to rotate
t = torch.randn(batch_size, heads, seq_len, feature_dim)

# Apply learned rotations
t_rotated = apply_learned_rotations(rotations, t)
```

--------------------------------

### Initialize RotaryEmbedding with Configuration Object

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/types.md

Use a dictionary to pass configuration parameters to the RotaryEmbedding constructor.

```python
config = {
    'dim': 64,
    'use_xpos': True,
    'xpos_scale_base': 512,
    'interpolate_factor': 1.0,
}
rotary_emb = RotaryEmbedding(**config)
```

--------------------------------

### Applying Scaling and Positional Offsets

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/apply-rotary-emb.md

Applies custom scaling factors or positional offsets, useful for XPos or caching scenarios.

```python
# Apply with custom scaling (used by XPos internally)
scale = 0.95  # Example: slight dampening
q_rot = apply_rotary_emb(freqs, q, scale=scale)

# With positional offset (for cache scenarios)
q_rot = apply_rotary_emb(freqs, q, scale=scale, seq_dim=-2)
```

--------------------------------

### rotate_queries_and_keys

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/rotary-embedding.md

Rotate queries and keys with XPos scaling for length-extrapolatable rotary embeddings. This method requires use_xpos=True.

```APIDOC
## rotate_queries_and_keys(q, k, seq_dim=None)

### Description
Rotate queries and keys with XPos scaling for length-extrapolatable rotary embeddings. Only works when use_xpos=True.

### Parameters
- **q** (Tensor) - Required - Query tensor. Shape: (..., seq_len, feature_dim)
- **k** (Tensor) - Required - Key tensor. Must have same shape as q for XPos
- **seq_dim** (int | None) - Optional - Sequence dimension. Defaults to -2

### Returns
Tuple of (rotated_queries, rotated_keys) with same shapes as inputs.
```

--------------------------------

### Import flash_attn_with_rotary

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/flash-attn-with-rotary.md

Import the function from the rotary_embedding_torch package.

```python
from rotary_embedding_torch.flash_attn_with_rotary import flash_attn_with_rotary
```

--------------------------------

### Initialize frequency bases

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/helpers-and-internals.md

Initializes frequency bases used for position multiplication. Choose between geometric, linear, or constant modes based on the desired frequency distribution.

```python
freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
```

```python
freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
```

```python
freqs = torch.ones(num_freqs).float()
```

--------------------------------

### Retrieve XPos Scaling Factors

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/rotary-embedding.md

Calculates exponential decay scaling factors for length extrapolation. Only functional when use_xpos=True is enabled.

```python
rotary_emb = RotaryEmbedding(dim=64, use_xpos=True, xpos_scale_base=512)
positions = torch.arange(1024)
scales = rotary_emb.get_scale(positions, seq_len=1024)
print(scales.shape)  # torch.Size([1024, 64])
```

--------------------------------

### Initialize RotaryEmbedding for Language Models

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/configuration.md

Standard configuration for language modeling tasks with default parameters.

```python
rotary_emb = RotaryEmbedding(
    dim=64,
    freqs_for='lang',
    theta=10000,
    use_xpos=False,
    interpolate_factor=1.0,
    cache_max_seq_len=8192
)
```

--------------------------------

### Rotate queries and keys with KV cache

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/rotary-embedding.md

Handles rotation for inference scenarios where keys are retrieved from a KV cache. Automatically manages position offsets and supports XPos scaling.

```python
import torch
from rotary_embedding_torch import RotaryEmbedding

rotary_emb = RotaryEmbedding(dim=64)

# At inference, single query token, full KV cache
q = torch.randn(1, 8, 1, 64)      # Current token
k_cache = torch.randn(1, 8, 512, 64)  # Accumulated cache

q_rot, k_rot = rotary_emb.rotate_queries_with_cached_keys(q, k_cache)
# q_rot.shape: (1, 8, 1, 64)
# k_rot.shape: (1, 8, 512, 64)
```

--------------------------------

### Initialize RotaryEmbedding with Learnable Frequencies

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/configuration.md

Configuration using learnable frequency parameters, which requires disabling caching.

```python
rotary_emb = RotaryEmbedding(
    dim=64,
    learned_freq=True,
    cache_if_possible=False,  # Disable caching for learned params
    use_xpos=False
)
```

--------------------------------

### Parameter Table Format

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/README.md

Standard markdown table format for documenting function parameters.

```markdown
| Parameter | Type | Required | Default | Description |
|-----------|------|----------|---------|-------------|
| name | type | ✓ or ✗ | value | description |
```

--------------------------------

### Broadcasting Learned Rotations with Frequency Ranges

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/apply-learned-rotations.md

Demonstrates how learned rotation parameters are broadcasted against frequency ranges to produce the final rotation matrix.

```python
rotations * freq_ranges  # Broadcasting applies

# If rotations: (num_rot,)
# And freq_ranges: (num_rot, num_freqs)
# Result: (num_rot, num_freqs)
```

--------------------------------

### Initialize RotaryEmbedding

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/rotary-embedding.md

Constructor signature for configuring the embedding layer, including parameters for dimension, frequency modes, and scaling options.

```python
def __init__(
    self,
    dim: int,
    custom_freqs: Tensor | None = None,
    freqs_for: Literal['lang', 'pixel', 'constant'] = 'lang',
    theta: int = 10000,
    max_freq: int = 10,
    num_freqs: int = 1,
    learned_freq: bool = False,
    use_xpos: bool = False,
    xpos_scale_base: int = 512,
    interpolate_factor: float = 1.,
    theta_rescale_factor: float = 1.,
    seq_before_head_dim: bool = False,
    cache_if_possible: bool = True,
    cache_max_seq_len: int = 8192
) -> None
```

--------------------------------

### get_scale

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/rotary-embedding.md

Compute XPos scaling factors for length extrapolation. Only available when use_xpos=True.

```APIDOC
## get_scale(t, seq_len=None, offset=0)

### Description
Compute XPos scaling factors for length extrapolation. Only available when use_xpos=True.

### Parameters
- **t** (Tensor) - Required - Sequence position indices where scales are needed
- **seq_len** (int | None) - Optional - Explicit sequence length for caching optimization
- **offset** (int) - Optional - Position offset when caching. Defaults to 0

### Returns
Tensor with shape (seq_len, dim) containing per-position scales.
```

--------------------------------

### Import RotaryEmbedding

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/rotary-embedding.md

Import the main class from the library.

```python
from rotary_embedding_torch import RotaryEmbedding
```

--------------------------------

### Define default() utility

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/helpers-and-internals.md

Provides a default value if the input value is None.

```python
def default(val: T, d: T) -> T:
    return val if exists(val) else d
```

```python
seq_dim = default(seq_dim, self.default_seq_dim)
device = default(device, self.device)
dtype = default(dtype, self.cached_freqs.dtype)

# Instead of:
seq_dim = seq_dim if seq_dim is not None else self.default_seq_dim
```

--------------------------------

### Configure Triton Kernel Grid

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/helpers-and-internals.md

Defines the compute grid for Triton kernel execution in Flash Attention.

```python
grid = lambda META: (
    triton.cdiv(seq_len_q, META['BLOCK_M']),
    batch * heads_q,
    1
)
```

--------------------------------

### Configure Sequence Interpolation

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/configuration.md

Rescales positions to extend context length during fine-tuning.

```python
# Pretraining
rotary_emb = RotaryEmbedding(dim=64, interpolate_factor=1.0)

# Fine-tuning for 4K context (2× pretraining length)
rotary_emb = RotaryEmbedding(dim=64, interpolate_factor=2.0)

# Fine-tuning for 8K context (4× pretraining length)
rotary_emb = RotaryEmbedding(dim=64, interpolate_factor=4.0)
```

--------------------------------

### Configure Learnable Frequencies

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/configuration.md

Enables gradient-based optimization of frequency parameters during training.

```python
# Learnable frequencies
rotary_emb = RotaryEmbedding(dim=64, learned_freq=True)

# Frequencies become parameters with gradients
print(rotary_emb.freqs.requires_grad)  # True

# Can be optimized with main model
optimizer.add_param_group({'params': [rotary_emb.freqs]})
```

--------------------------------

### Configure Theta Rescale Factor

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/configuration.md

Apply NTK-aware scaling to extend context length without fine-tuning.

```python
import math

# Extend from 4K to 32K tokens (8x scaling)
rotary_emb = RotaryEmbedding(
    dim=64,
    theta=10000,
    theta_rescale_factor=math.sqrt(8)  # ≈ 2.83
)

# This is equivalent to much larger theta value but applied uniformly
```

--------------------------------

### Rotate Queries and Keys with XPos

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/rotary-embedding.md

Applies length-extrapolatable rotary embeddings to query and key tensors. Requires the RotaryEmbedding instance to be initialized with use_xpos=True.

```python
import torch
from rotary_embedding_torch import RotaryEmbedding

rotary_emb = RotaryEmbedding(dim=64, use_xpos=True)

q = torch.randn(2, 8, 512, 64)
k = torch.randn(2, 8, 512, 64)

# Apply XPos rotations
q_rot, k_rot = rotary_emb.rotate_queries_and_keys(q, k)
```

--------------------------------

### Force Reference Implementation for Triton

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/index.md

Use the reference implementation for debugging when the Triton kernel is unavailable or failing.

```python
# To force reference implementation for debugging:
from rotary_embedding_torch.flash_attn_with_rotary import get_flash_attention_fused
flash_attn_ref = get_flash_attention_fused(force_reference=True)
output = flash_attn_ref(q, k, v, rotary_pos_emb=freqs)
```

--------------------------------

### broadcat()

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/index.md

Broadcasting and concatenation utility.

```APIDOC
## broadcat(tensors, dim=-1)

### Description
Broadcasting and concatenation utility.

### Signature
`broadcat(tensors: list[Tensor], dim: int = -1) -> Tensor`
```

--------------------------------

### Import apply_rotary_emb

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/apply-rotary-emb.md

Import the function from the rotary_embedding_torch package.

```python
from rotary_embedding_torch import apply_rotary_emb
```

--------------------------------

### Configure Frequency Computation Modes

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/configuration.md

Select the computation mode for base frequencies based on the target domain.

```python
# Language model (default)
rotary_emb = RotaryEmbedding(dim=64, freqs_for='lang')

# Vision transformer
rotary_emb = RotaryEmbedding(dim=48, freqs_for='pixel', max_freq=256)

# Constant frequencies
rotary_emb = RotaryEmbedding(dim=64, freqs_for='constant', num_freqs=8)
```

--------------------------------

### Configure RotaryEmbedding Caching

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/configuration.md

Toggle frequency caching to balance speed and memory usage.

```python
# With caching (default)
rotary_emb = RotaryEmbedding(dim=64, cache_if_possible=True)

# Without caching (for memory-constrained scenarios)
rotary_emb = RotaryEmbedding(dim=64, cache_if_possible=False)
```

--------------------------------

### RotaryEmbedding Configuration Presets

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/index.md

Presets for different model architectures and sequence length requirements.

```python
RotaryEmbedding(dim=64)
```

```python
RotaryEmbedding(dim=64, theta=100000, cache_max_seq_len=32768)
```

```python
RotaryEmbedding(dim=64, use_xpos=True, theta=100000)
```

```python
RotaryEmbedding(dim=48, freqs_for='pixel', max_freq=256)
```

--------------------------------

### Configure XPos Length Extrapolatable Embeddings

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/README.md

Enabling XPos for autoregressive transformers to improve extrapolation to sequence lengths beyond training.

```python
import torch
from rotary_embedding_torch import RotaryEmbedding

# instantiate the positional embedding in your transformer and pass to all your attention layers

rotary_emb = RotaryEmbedding(
    dim = 32,
    use_xpos = True   # set this to True to make rotary embeddings extrapolate better to sequence lengths greater than the one used at training time
)

# mock queries and keys - dimensions should end with (seq_len, feature dimension), and any number of preceding dimensions (batch, heads, etc)

q = torch.randn(1, 8, 1024, 64) # queries - (batch, heads, seq len, dimension of head)
k = torch.randn(1, 8, 1024, 64) # keys

# apply the rotations to your queries and keys after the heads have been split out, but prior to the dot product and subsequent softmax (attention)

# instead of using `rotate_queries_or_keys`, you will use `rotate_queries_and_keys`, the rest is taken care of

q, k = rotary_emb.rotate_queries_and_keys(q, k)
```

--------------------------------

### Import apply_learned_rotations

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/apply-learned-rotations.md

Import the function from the rotary_embedding_torch package.

```python
from rotary_embedding_torch import apply_learned_rotations
```

--------------------------------

### Initialize RotaryEmbedding with Extrapolation

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/configuration.md

Configuration for length extrapolation, enabling inference on sequences longer than pretraining lengths.

```python
rotary_emb = RotaryEmbedding(
    dim=64,
    freqs_for='lang',
    theta=100000,  # Extended theta
    theta_rescale_factor=1.0,
    use_xpos=True,  # Length-extrapolatable
    xpos_scale_base=512,
    cache_max_seq_len=32768  # Larger cache
)
```

--------------------------------

### Cite XPos Research

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/index.md

BibTeX citation for the XPos length-extrapolatable transformer paper.

```bibtex
@article{sun2022length,
    title={A Length-Extrapolatable Transformer},
    author={Sun, Yutao and Dong, Li and Patra, Barun and Ma, Shuming and Huang, Shaohan and Benhaim, Alon and Chaudhary, Vishrav and Song, Xia and Wei, Furu},
    year={2022}
}
```

--------------------------------

### Internal Helper Functions

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/types.md

Utility functions used internally for checking existence and providing default values.

```python
# In rotary_embedding_torch.py
def exists(val) -> bool:
    return val is not None

def default(val: T, d: T) -> T:
    return val if exists(val) else d
```

--------------------------------

### Length Extrapolatable Embeddings (XPos)

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/index.md

Enables XPos for length extrapolation during training and inference.

```python
# Training with XPos enabled
rotary_emb = RotaryEmbedding(dim=64, use_xpos=True)

# During training and inference
q = torch.randn(batch, heads, seq_len, 64)
k = torch.randn(batch, heads, seq_len, 64)

q_rot, k_rot = rotary_emb.rotate_queries_and_keys(q, k)
```

--------------------------------

### Providing Cosine and Sine Directly

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/flash-attn-with-rotary.md

Manually provide precomputed cos and sin tensors instead of the rotary embedding frequencies.

```python
import torch
from rotary_embedding_torch import RotaryEmbedding
from rotary_embedding_torch.flash_attn_with_rotary import flash_attn_with_rotary

rotary_emb = RotaryEmbedding(dim=64)
seq_len = 1024
freqs = rotary_emb(torch.arange(seq_len)).cuda()

# Manually compute cos/sin in the expected shape
cos = torch.cos(freqs).unsqueeze(0).unsqueeze(0).unsqueeze(0).unsqueeze(0)  # (1,1,1,1,seq_len,dim)
sin = torch.sin(freqs).unsqueeze(0).unsqueeze(0).unsqueeze(0).unsqueeze(0)

q = torch.randn(2, 8, seq_len, 64).cuda()
k = torch.randn(2, 8, seq_len, 64).cuda()
v = torch.randn(2, 8, seq_len, 64).cuda()

output = flash_attn_with_rotary(
    q, k, v,
    cos=cos,
    sin=sin,
    is_causal=True
)
```

--------------------------------

### Apply Type Hints to Rotary Embedding Usage

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/types.md

Recommended pattern for type-hinting tensors and class instances when applying rotations.

```python
from rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb
import torch

def apply_rotations(q: torch.Tensor, k: torch.Tensor, dim: int) -> tuple[torch.Tensor, torch.Tensor]:
    rotary_emb: RotaryEmbedding = RotaryEmbedding(dim=dim)
    q_rot: torch.Tensor = rotary_emb.rotate_queries_or_keys(q)
    k_rot: torch.Tensor = rotary_emb.rotate_queries_or_keys(k)
    return q_rot, k_rot
```

--------------------------------

### BibTeX Citation for NTK-Aware Scaled RoPE

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/README.md

Citation for the community-developed NTK-aware scaling method for RoPE.

```bibtex
@misc{bloc97-2023
    title   = {NTK-Aware Scaled RoPE allows LLaMA models to have extended (8k+) context size without any fine-tuning and minimal perplexity degradation.},
    author  = {/u/bloc97},
    url     = {https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/}
}
```

--------------------------------

### Define Optional types

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/types.md

Indicates parameters that accept either a specific type or None.

```python
Tensor | None
int | None
float | None
```

--------------------------------

### Configure XPos Decay Constant

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/configuration.md

Adjusts the decay rate of position-dependent scales using the xpos_scale_base parameter.

```python
# Slower decay (gentler extrapolation)
rotary_emb = RotaryEmbedding(
    dim=64,
    use_xpos=True,
    xpos_scale_base=1024  # Larger decay constant
)

# Faster decay (stronger position signal)
rotary_emb = RotaryEmbedding(
    dim=64,
    use_xpos=True,
    xpos_scale_base=256  # Smaller decay constant
)
```

--------------------------------

### Access Device Property

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/helpers-and-internals.md

Retrieves the device of the embedding using a registered dummy buffer to avoid direct data access.

```python
@property
def device(self) -> torch.device:
    return self.dummy.device
```

--------------------------------

### Initialize RotaryEmbedding with Interpolation

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/README.md

Configure the embedding layer with an interpolation factor to extend context length for pretrained models.

```python
import torch
from rotary_embedding_torch import RotaryEmbedding

rotary_emb = RotaryEmbedding(
    dim = 32,
    interpolate_factor = 2.    # add this line of code to pretrained model and fine-tune for ~1000 steps, as shown in paper
)
```

--------------------------------

### Function Signature for apply_learned_rotations

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/apply-learned-rotations.md

The definition of the function parameters and return type.

```python
def apply_learned_rotations(
    rotations: Tensor,
    t: Tensor,
    start_index: int = 0,
    freq_ranges: Tensor | None = None
) -> Tensor
```

--------------------------------

### Initialize RotaryEmbedding for Vision Transformers

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/configuration.md

Configuration optimized for pixel-based inputs such as video data.

```python
rotary_emb = RotaryEmbedding(
    dim=48,
    freqs_for='pixel',
    max_freq=256,
    use_xpos=False,
    cache_max_seq_len=2048
)
```

--------------------------------

### BibTeX Citation for Positional Interpolation

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/README.md

Citation for the paper on extending context windows via positional interpolation.

```bibtex
@inproceedings{Chen2023ExtendingCW,
    title   = {Extending Context Window of Large Language Models via Positional Interpolation},
    author  = {Shouyuan Chen and Sherman Wong and Liangjian Chen and Yuandong Tian},
    year      = {2023}
}
```

--------------------------------

### apply_learned_rotations()

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/index.md

Apply learnable rotation parameters instead of fixed frequencies.

```APIDOC
## apply_learned_rotations(rotations, t, start_index=0, freq_ranges=None)

### Description
Apply learnable rotation parameters instead of fixed frequencies.

### Signature
`apply_learned_rotations(rotations: Tensor, t: Tensor, start_index: int = 0, freq_ranges: Tensor | None = None) -> Tensor`
```

--------------------------------

### Import public modules from rotary_embedding_torch

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/types.md

Use these stable exports to access core functionality. These functions and classes are guaranteed to remain stable across versions.

```python
from rotary_embedding_torch.rotary_embedding_torch import (
    apply_rotary_emb,
    RotaryEmbedding,
    apply_learned_rotations,
    broadcat
)
```

--------------------------------

### Configure Tensor Layout

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/configuration.md

Sets the dimension ordering for input tensors to match specific model architectures.

```python
# Standard multi-head attention layout
rotary_emb = RotaryEmbedding(dim=64, seq_before_head_dim=False)
q = torch.randn(batch, heads, seq_len, 64)
q_rot = rotary_emb.rotate_queries_or_keys(q)

# Alternative layout
rotary_emb = RotaryEmbedding(dim=64, seq_before_head_dim=True)
q = torch.randn(batch, seq_len, heads, 64)
q_rot = rotary_emb.rotate_queries_or_keys(q)
```

--------------------------------

### Compute Per-Position Scaling

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/helpers-and-internals.md

Generates a position-dependent decay envelope by centering indices and applying exponential scaling.

```python
power = (t - len(t) // 2) / self.scale_base
scale = self.scale ** rearrange(power, 'n -> n 1')
scale = repeat(scale, 'n d -> n (d r)', r=2)
```

--------------------------------

### Module Imports for Type Annotations

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/types.md

Required imports for utilizing PyTorch and typing features within the module.

```python
from torch import nn, einsum, broadcast_tensors, is_tensor, tensor, Tensor
from typing import Literal
```

--------------------------------

### Function Return Type Signatures

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/types.md

Common return type patterns for library functions, including single Tensor and tuple outputs.

```python
def forward(...) -> Tensor
def apply_rotary_emb(...) -> Tensor
def apply_learned_rotations(...) -> Tensor
```

```python
def rotate_queries_with_cached_keys(...) -> tuple[Tensor, Tensor]
def rotate_queries_and_keys(...) -> tuple[Tensor, Tensor]
```

--------------------------------

### Mathematical Operation for Learned Rotations

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/apply-learned-rotations.md

Conceptual representation of the rotation application logic.

```text
if freq_ranges provided:
    weighted_rotations = einsum('..., f -> ... f', rotations, freq_ranges)
    weighted_rotations = rearrange to pair format

rotation_matrix * exp(i * learned_angle)
```

--------------------------------

### RotaryEmbedding Initialization

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/configuration.md

The RotaryEmbedding class is initialized with various parameters to control frequency generation, context length scaling, and mode-specific behavior.

```APIDOC
## RotaryEmbedding(dim, custom_freqs=None, freqs_for='lang', theta=10000, theta_rescale_factor=1.0, max_freq=10, num_freqs=1)

### Description
Initializes the RotaryEmbedding module with specified frequency generation parameters.

### Parameters
- **dim** (int) - Required - The number of dimensions to embed rotationally. Must be an even integer > 0.
- **custom_freqs** (Tensor | None) - Optional - Custom frequency tensor of shape (dim // 2,) to use instead of computed frequencies.
- **freqs_for** (Literal['lang', 'pixel', 'constant']) - Optional - Determines how base frequencies are computed. Defaults to 'lang'.
- **theta** (int) - Optional - Base for geometric frequency progression in 'lang' mode. Defaults to 10000.
- **theta_rescale_factor** (float) - Optional - NTK-aware scaling factor for adjusting theta. Defaults to 1.0.
- **max_freq** (int) - Optional - Maximum frequency for 'pixel' mode. Defaults to 10.
- **num_freqs** (int) - Optional - Number of constant frequencies for 'constant' mode. Defaults to 1.

### Example
```python
from rotary_embedding_torch import RotaryEmbedding

# Standard language model initialization
rotary_emb = RotaryEmbedding(dim=64, theta=10000)

# Vision transformer initialization
rotary_emb = RotaryEmbedding(dim=48, freqs_for='pixel', max_freq=256)
```
```

--------------------------------

### Multi-Head Attention with GQA

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/flash-attn-with-rotary.md

Supports Grouped Query Attention (GQA) where KV heads are fewer than query heads, automatically expanding them.

```python
import torch
from rotary_embedding_torch import RotaryEmbedding
from rotary_embedding_torch.flash_attn_with_rotary import flash_attn_with_rotary

rotary_emb = RotaryEmbedding(dim=64)
seq_len = 1024
freqs = rotary_emb(torch.arange(seq_len)).cuda()

# Grouped query attention: 8 query heads, 2 KV heads
q = torch.randn(2, 8, seq_len, 64).cuda()
k = torch.randn(2, 2, seq_len, 64).cuda()  # Fewer heads
v = torch.randn(2, 2, seq_len, 64).cuda()

# Automatic expansion: each KV head serves 4 query heads
output = flash_attn_with_rotary(
    q, k, v,
    rotary_pos_emb=freqs,
    is_causal=True
)

print(output.shape)  # torch.Size([2, 8, 1024, 64])
```

--------------------------------

### Validate Cache Eligibility

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/helpers-and-internals.md

Determines if frequency caching is permissible based on configuration, sequence length, and mode constraints.

```python
should_cache = (
    self.cache_if_possible and
    not self.learned_freq and
    exists(seq_len) and
    self.freqs_for != 'pixel' and
    (offset + seq_len) <= self.cache_max_seq_len
)
```

--------------------------------

### Specify Tensor Dtypes Explicitly

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/types.md

Best practice for ensuring tensor precision by defining dtypes or device placement explicitly.

```python
# Good: explicit dtype
q = torch.randn(batch, heads, seq_len, head_dim, dtype=torch.float32, device='cuda')

# Also OK: infer from model
q = torch.randn(batch, heads, seq_len, head_dim).to(model.device)
```

--------------------------------

### Define exists() utility

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/helpers-and-internals.md

A type-safe utility to check if a value is not None.

```python
def exists(val: Any) -> bool:
    return val is not None
```

```python
if exists(custom_freqs):
    freqs = custom_freqs
else:
    freqs = compute_freqs()

# Instead of:
if custom_freqs is not None:
    freqs = custom_freqs
```

--------------------------------

### device

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/rotary-embedding.md

Property to retrieve the device of the embedding buffers.

```APIDOC
## device

### Description
Returns the device of the embedding buffers. Used to ensure frequency tensors are on the correct device.

### Signature
`@property def device(self) -> torch.device`
```

--------------------------------

### Train with Learned Rotations in Attention

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/apply-learned-rotations.md

Integrates learned rotation parameters directly into a custom nn.Module for use within an attention mechanism.

```python
import torch
import torch.nn as nn
from rotary_embedding_torch import apply_learned_rotations

class LearnedRotaryAttention(nn.Module):
    def __init__(self, dim_head, num_learned_rotations=16):
        super().__init__()
        # Learnable rotation parameters
        self.rotation_params = nn.Parameter(
            torch.zeros(num_learned_rotations)
        )
        
    def forward(self, q, k, v):
        # Apply learned rotations
        q = apply_learned_rotations(self.rotation_params, q)
        k = apply_learned_rotations(self.rotation_params, k)
        
        # Continue with attention
        sim = torch.einsum('b h i d, b h j d -> b h i j', q, k)
        attn = sim.softmax(dim=-1)
        return torch.einsum('b h i j, b h j d -> b h i d', attn, v)

# Usage
attention = LearnedRotaryAttention(dim_head=64)
q = torch.randn(2, 8, 512, 64)
k = torch.randn(2, 8, 512, 64)
v = torch.randn(2, 8, 512, 64)

output = attention(q, k, v)
```

--------------------------------

### Concatenate Tensors on Different Dimensions

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/broadcat.md

Demonstrates broadcasting and concatenation on the first dimension.

```python
# Concatenate on dim=0 (first dimension)
t1 = torch.randn(1, 32)
t2 = torch.randn(8, 32)
result = broadcat([t1, t2], dim=0)
# Both broadcast to (8, 32)
# Concatenate on dim=0: (16, 32)
```

--------------------------------

### Construct frequency matrix

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/helpers-and-internals.md

Computes the frequency matrix by multiplying position indices with frequency bases and repeating values to match the expected cos/sin pair format.

```python
freqs = einsum('..., f -> ... f', t.type(freqs.dtype), freqs)
freqs = repeat(freqs, '... n -> ... (n r)', r=2)
```

--------------------------------

### Partial Rotation with start_index

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/apply-rotary-emb.md

Applies rotary embeddings to a subset of feature dimensions, leaving the remaining features unchanged.

```python
# If using both rotary embeddings and other positional encodings
# Rotate only the first 32 dimensions
dim_rotary = 32
dim_other = 32

q = torch.randn(batch, heads, seq_len, dim_rotary + dim_other)

# Only the first 32 features are rotated
freqs = rotary_emb(torch.arange(seq_len))  # dim=32
q_rot = apply_rotary_emb(freqs, q, start_index=0)
# Features [32:] remain unchanged
```

--------------------------------

### Apply Rotary Embeddings

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/README.md

Instantiate the embedding layer and apply rotations to query and key tensors before the attention mechanism.

```python
import torch
from rotary_embedding_torch import RotaryEmbedding

# instantiate the positional embedding in your transformer and pass to all your attention layers

rotary_emb = RotaryEmbedding(dim = 32)

# mock queries and keys - dimensions should end with (seq_len, feature dimension), and any number of preceding dimensions (batch, heads, etc)

q = torch.randn(1, 8, 1024, 64) # queries - (batch, heads, seq len, dimension of head)
k = torch.randn(1, 8, 1024, 64) # keys

# apply the rotations to your queries and keys after the heads have been split out, but prior to the dot product and subsequent softmax (attention)

q = rotary_emb.rotate_queries_or_keys(q)
k = rotary_emb.rotate_queries_or_keys(k)

# then do your attention with your queries (q) and keys (k) as usual
```

--------------------------------

### Execute multi-dimensional broadcasting

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/broadcat.md

Handles tensors with varying dimensions by broadcasting them to a common shape before concatenation.

```python
import torch
from rotary_embedding_torch import broadcat

# Different batch sizes, same features
t1 = torch.randn(1, 4, 32)    # (1, 4, 32)
t2 = torch.randn(8, 1, 32)    # (8, 1, 32)
t3 = torch.randn(8, 4, 32)    # (8, 4, 32)

# Broadcast all to (8, 4, 32), then concatenate
result = broadcat([t1, t2, t3], dim=-1)
# All broadcast to (8, 4, 32)
# Concatenate along last dim: (8, 4, 96)
print(result.shape)  # torch.Size([8, 4, 96])
```

--------------------------------

### Combine Hybrid Features with Broadcasting

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/broadcat.md

Concatenate fixed and learned features where the learned parameters have a batch dimension of 1 and require broadcasting.

```python
# When some features come from fixed components and others from learned components
# but they have different batch dimension handling

import torch
from rotary_embedding_torch import broadcat

def hybrid_features(x, learned_params):
    """
    x: shape (batch, seq, d_fixed)
    learned_params: shape (1, seq, d_learned) - broadcast across batch
    """
    # learned_params will be broadcast to (batch, seq, d_learned)
    features = broadcat([x, learned_params], dim=-1)
    return features  # (batch, seq, d_fixed + d_learned)

batch_size = 32
seq_len = 512
x = torch.randn(batch_size, seq_len, 256)
learned = torch.randn(1, seq_len, 128)

features = hybrid_features(x, learned)
print(features.shape)  # torch.Size([32, 512, 384])
```

--------------------------------

### Perform basic broadcasting and concatenation

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/broadcat.md

Broadcasts tensors with compatible shapes and concatenates them along the last dimension.

```python
import torch
from rotary_embedding_torch import broadcat

# Tensors with broadcastable shapes
t1 = torch.randn(1, 32)  # shape: (1, 32)
t2 = torch.randn(8, 32)  # shape: (8, 32)

# Broadcast and concatenate along dim=-1
result = broadcat([t1, t2], dim=-1)
# Broadcasting: both become (8, 32)
# Concatenation along dim=-1: (8, 64)
print(result.shape)  # torch.Size([8, 64])
```

--------------------------------

### Apply Learned Rotations with Frequency Range Weighting

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/apply-learned-rotations.md

Applies rotations with an additional frequency modulation layer, requiring the input tensor dimension to be at least twice the number of rotation parameters.

```python
import torch
from rotary_embedding_torch import apply_learned_rotations

# Learned rotation angles
num_rotations = 16
rotations = torch.nn.Parameter(torch.randn(num_rotations))

# Frequency range weights (can come from another learned layer)
# Maps rotation parameters to frequency bins
freq_ranges = torch.randn(num_rotations, 8)  # 8 frequency dimensions

# Tensor
t = torch.randn(batch, heads, seq_len, 128)  # Must be >= num_rotations * 2

# Apply with frequency modulation
t_rotated = apply_learned_rotations(
    rotations, 
    t, 
    freq_ranges=freq_ranges
)
```