PoPE PyTorch

https://github.com/hunterhogan/pope-pytorch
Admin
PoPE PyTorch is an efficient implementation of polar coordinate positional embeddings for...
Tokens:7,063
Snippets:17
Trust Score:8.3
Update:6 months ago
Show doc for...
Tokens:
Raw
### Install PoPE-pytorch

Source: https://github.com/hunterhogan/pope-pytorch/blob/main/README.md

Install the library using pip.

```shell
pip install PoPE-pytorch
```

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

### Initialize and Use PoPE

Source: https://github.com/hunterhogan/pope-pytorch/blob/main/README.md

Initialize PoPE with dimension and heads, then apply to queries and keys for attention.

```python
import torch
from PoPE_pytorch import PoPE

# define pope

pope = PoPE(64, heads = 8)

# pass in sequence length

pos_emb = pope(1024)

# queries and keys in attention

q = torch.randn(1, 8, 1024, 64)
k = torch.randn(1, 8, 1024, 64)

# training

rotated_q, rotated_k = pope.apply_pope_to_qk(pos_emb, q, k)

# inference

rotated_q, rotated_k = pope.apply_pope_to_qk(pos_emb, q[..., -1:, :], k)
```

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

### Initialize and Use AxialPoPE for Video

Source: https://github.com/hunterhogan/pope-pytorch/blob/main/README.md

Initialize AxialPoPE for video data, splitting the dimension across time and spatial dimensions.

```python
# for video (e.g. 8 frames, 16x16 frames)
# split 96 dim into 32 (t), 32 (x), 32 (y)

pope_video = AxialPoPE(
    dim = 96,
    heads = 8,
    axial_dims = (32, 32, 32)
)

pos_emb_video = pope_video((8, 16, 16)) # (2048, 96) frequencies

# queries and keys
# then apply to q, k as usual

q = torch.randn(1, 8, 2048, 96)
k = torch.randn(1, 8, 2048, 96)

rotated_q, rotated_k = AxialPoPE.apply_pope_to_qk(pos_emb_video, q, k)
```

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

### Initialize and Use AxialPoPE for Images

Source: https://github.com/hunterhogan/pope-pytorch/blob/main/README.md

Initialize AxialPoPE for image data, splitting the dimension across specified axial dimensions.

```python
import torch
from PoPE_pytorch import AxialPoPE

# axial pope for images (e.g. 16x16)
# split 64 dim into 32 (x) and 32 (y)

pope = AxialPoPE(
    dim = 64,
    heads = 8,
    axial_dims = (32, 32)
)

pos_emb = pope((16, 16)) # (256, 64) frequencies
```

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

### Fused Flash Attention with PoPE

Source: https://github.com/hunterhogan/pope-pytorch/blob/main/README.md

Utilize PoPE within a fused Flash Attention implementation for efficiency.

```python
import torch
from PoPE_pytorch import PoPE, flash_attn_with_pope

# pope

pope = PoPE(dim = 32, heads = 8).cuda()

# queries, keys, values for attention

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

pos_emb = pope(1024)

mask = torch.ones((2, 1024)).bool().cuda()

out = flash_attn_with_pope(q, k, v, pos_emb = pos_emb, causal = True, mask = mask)

assert out.shape == (2, 8, 1024, 64)
```

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

### Instantiate and Apply PoPE for 1D Sequences

Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt

Instantiate the PoPE module and generate positional embeddings for a given sequence length. Apply these embeddings to query and key tensors for training or inference.

```python
import torch
from PoPE_pytorch import PoPE

# --- Instantiate ---
# dim: per-head rotation dimension (must be <= head_dim in Q/K)
# heads: number of attention heads
pope = PoPE(dim=64, heads=8, theta=10000)

# --- Generate positional embeddings for a sequence of length 1024 ---
pos_emb = pope(1024)
# pos_emb.freqs  shape: (1024, 64)
# pos_emb.bias   shape: (8, 64)

# --- Apply to Q and K (training: full context) ---
q = torch.randn(1, 8, 1024, 64)
k = torch.randn(1, 8, 1024, 64)

rotated_q, rotated_k = pope.apply_pope_to_qk(pos_emb, q, k)
assert rotated_q.shape == (1, 8, 1024, 128)  # dim doubles after polar interleaving
assert rotated_k.shape == (1, 8, 1024, 128)

# --- Apply to Q and K (inference: single new token, cached keys) ---
rotated_q_inf, rotated_k_inf = pope.apply_pope_to_qk(pos_emb, q[..., -1:, :], k)
assert rotated_q_inf.shape == (1, 8, 1,    128)
assert rotated_k_inf.shape == (1, 8, 1024, 128)

# --- Use explicit position tensor (e.g. for packed/offset sequences) ---
positions = torch.arange(512, 512 + 1024).float()
pos_emb_offset = pope(positions)
```

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

### Instantiate and Apply AxialPoPE for Multi-Dimensional Grids

Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt

Instantiate AxialPoPE for multi-dimensional grids like images or video. Generate positional embeddings using grid dimensions and apply them to query and key tensors.

```python
import torch
from PoPE_pytorch import AxialPoPE

# --- 2D image grid (16×16 patches, head_dim=64 split as 32+32) ---
pope_2d = AxialPoPE(dim=64, heads=8, axial_dims=(32, 32))

pos_emb_2d = pope_2d((16, 16))          # auto meshgrid → (256, 64) freqs
assert pos_emb_2d.freqs.shape == (256, 64)

q = torch.randn(1, 8, 256, 64)
k = torch.randn(1, 8, 256, 64)
rotated_q, rotated_k = AxialPoPE.apply_pope_to_qk(pos_emb_2d, q, k)
assert rotated_q.shape == (1, 8, 256, 128)
```

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

### 3D Positional Embeddings with AxialPoPE

Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt

Demonstrates the creation and application of 3D positional embeddings using AxialPoPE for video data. Shows both automatic meshgrid generation and manual position definition.

```python
pope_3d = AxialPoPE(dim=96, heads=8, axial_dims=(32, 32, 32))

pos_emb_3d = pope_3d((8, 16, 16)))       # auto meshgrid → (2048, 96) freqs
assert pos_emb_3d.freqs.shape == (2048, 96)

q3 = torch.randn(1, 8, 2048, 96)
k3 = torch.randn(1, 8, 2048, 96)
rotated_q3, rotated_k3 = AxialPoPE.apply_pope_to_qk(pos_emb_3d, q3, k3)
assert rotated_q3.shape == (1, 8, 2048, 192)
```

```python
# Manual positions (e.g. irregular or cropped grids) \
custom_pos = AxialPoPE.get_grid_positions(4, 4, device='cpu')  # (16, 2) \
pos_emb_manual = pope_2d(custom_pos)
```

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

### Apply PoPE to QK Tensors with Customizations

Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt

Use the standalone `apply_pope_to_qk` function to apply positional embeddings. Supports partial rotation, custom magnitude activations, and complex tensor output.

```python
import torch
import torch.nn.functional as F
from PoPE_pytorch import PoPE
from PoPE_pytorch.pope import apply_pope_to_qk

pope = PoPE(dim=32, heads=4)  # only rotate the first 32 of 64 dims (partial)
pos_emb = pope(256)

q = torch.randn(2, 4, 256, 64)
k = torch.randn(2, 4, 256, 64)

# Partial rotation: first 32 dims rotated, remaining 32 passed through unchanged
rotated_q, rotated_k = apply_pope_to_qk(pos_emb, q, k)
assert rotated_q.shape == (2, 4, 256, 96)  # 32*2 + 32 = 96

# Using ReLU instead of softplus as magnitude activation
rotated_q2, rotated_k2 = apply_pope_to_qk(pos_emb, q, k, to_magnitude=F.relu)

# Return complex tensors instead of interleaved real
rotated_q_c, rotated_k_c = apply_pope_to_qk(pos_emb, q, k, return_complex=True)
assert rotated_q_c.is_complex()
```

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

### Flash Attention with PoPE Integration

Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt

A full attention function that fuses PoPE rotations into a Triton-based flash attention kernel or falls back to scaled_dot_product_attention. Supports various features like causal masking, GQA, and different head layouts.

```python
import torch
from PoPE_pytorch import PoPE, flash_attn_with_pope

pope = PoPE(dim=32, heads=8).cuda()
pos_emb = pope(1024)

# --- Basic causal self-attention (training) ---
q = torch.randn(2, 8, 1024, 64).cuda()
k = torch.randn(2, 8, 1024, 64).cuda()
v = torch.randn(2, 8, 1024, 64).cuda()
mask = torch.ones(2, 1024, dtype=torch.bool).cuda()

out = flash_attn_with_pope(q, k, v, pos_emb=pos_emb, causal=True, mask=mask)
assert out.shape == (2, 8, 1024, 64)
```

```python
# --- Inference: single new token, cached context (causal) ---
q_new = torch.randn(2, 8, 1, 64).cuda()
out_inf = flash_attn_with_pope(q_new, k, v, pos_emb=pos_emb, causal=True)
assert out_inf.shape == (2, 8, 1, 64)
```

```python
# --- Heads-last layout (b, n, h, d) ---
q_hl = torch.randn(2, 1024, 8, 64).cuda()
k_hl = torch.randn(2, 1024, 8, 64).cuda()
v_hl = torch.randn(2, 1024, 8, 64).cuda()
out_hl = flash_attn_with_pope(q_hl, k_hl, v_hl, pos_emb=pos_emb, causal=False,
                               head_dimension_at_first=False)
assert out_hl.shape == (2, 1024, 8, 64)
```

```python
# --- Force non-fused (CPU or debugging) ---
out_manual = flash_attn_with_pope(q, k, v, pos_emb=pos_emb, causal=True, fused=False)
```

```python
# --- Custom softmax scale + dropout ---
out_scaled = flash_attn_with_pope(q, k, v, pos_emb=pos_emb, causal=True,
                                   softmax_scale=0.125, dropout=0.1)
```

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

### flash_attn_with_pope

Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt

A complete multi-head self/cross-attention function that fuses PoPE rotations into a Triton-based flash attention kernel or falls back to `torch.nn.functional.scaled_dot_product_attention`. Supports various masking, dropout, GQA, and layouts.

```APIDOC
## flash_attn_with_pope

### Description
A complete multi-head self/cross-attention function that fuses PoPE rotations into a Triton-based flash attention kernel (when available on CUDA) or falls back to `torch.nn.functional.scaled_dot_product_attention`. Supports causal masking, key-padding masks, dropout, grouped query attention, and both heads-first `(b, h, n, d)` and heads-last `(b, n, h, d)` layouts.

### Method
`flash_attn_with_pope(q, k, v, pos_emb, causal=False, mask=None, softmax_scale=None, dropout=0.0, fused=True, head_dimension_at_first=True)`

### Parameters
- **q** (torch.Tensor) - Query tensor.
- **k** (torch.Tensor) - Key tensor.
- **v** (torch.Tensor) - Value tensor.
- **pos_emb** (torch.Tensor) - Positional embedding tensor.
- **causal** (bool, optional) - Whether to apply causal masking. Defaults to False.
- **mask** (torch.Tensor, optional) - Key-padding mask. Defaults to None.
- **softmax_scale** (float, optional) - Softmax scale factor. Defaults to None.
- **dropout** (float, optional) - Dropout rate for attention weights. Defaults to 0.0.
- **fused** (bool, optional) - Whether to use the fused kernel. Defaults to True.
- **head_dimension_at_first** (bool, optional) - Specifies if the head dimension is the first dimension after batch. Defaults to True.

### Request Example
```python
import torch
from PoPE_pytorch import PoPE, flash_attn_with_pope

pope = PoPE(dim=32, heads=8).cuda()
pos_emb = pope(1024)

# --- Basic causal self-attention (training) ---
q = torch.randn(2, 8, 1024, 64).cuda()
k = torch.randn(2, 8, 1024, 64).cuda()
v = torch.randn(2, 8, 1024, 64).cuda()
mask = torch.ones(2, 1024, dtype=torch.bool).cuda()

out = flash_attn_with_pope(q, k, v, pos_emb=pos_emb, causal=True, mask=mask)
assert out.shape == (2, 8, 1024, 64)

# --- Inference: single new token, cached context (causal) ---
q_new = torch.randn(2, 8, 1, 64).cuda()
out_inf = flash_attn_with_pope(q_new, k, v, pos_emb=pos_emb, causal=True)
assert out_inf.shape == (2, 8, 1, 64)

# --- Heads-last layout (b, n, h, d) ---
q_hl = torch.randn(2, 1024, 8, 64).cuda()
k_hl = torch.randn(2, 1024, 8, 64).cuda()
v_hl = torch.randn(2, 1024, 8, 64).cuda()
out_hl = flash_attn_with_pope(q_hl, k_hl, v_hl, pos_emb=pos_emb, causal=False,
                               head_dimension_at_first=False)
assert out_hl.shape == (2, 1024, 8, 64)

# --- Force non-fused (CPU or debugging) ---
out_manual = flash_attn_with_pope(q, k, v, pos_emb=pos_emb, causal=True, fused=False)

# --- Custom softmax scale + dropout ---
out_scaled = flash_attn_with_pope(q, k, v, pos_emb=pos_emb, causal=True,
                                   softmax_scale=0.125, dropout=0.1)
```

### Response
#### Success Response (200)
- **out** (torch.Tensor) - The output tensor after applying attention.
```

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

### AxialPoPE Module

Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt

`AxialPoPE` extends PoPE to multi-dimensional grids by splitting the feature dimension across independent axial embeddings. The `forward` method accepts position tensors or grid sizes to generate embeddings for multi-dimensional data.

```APIDOC
## AxialPoPE — Multi-Dimensional Positional Embedding

### Description
`AxialPoPE` extends PoPE to multi-dimensional grids by splitting the feature dimension across independent axial embeddings (e.g. x and y for images, t/x/y for video). The `forward` method accepts either an explicit `(N, num_axes)` position tensor or a tuple of grid sizes, which auto-generates a meshgrid.

### Usage (2D Image Grid)
```python
import torch
from PoPE_pytorch import AxialPoPE

# --- 2D image grid (16×16 patches, head_dim=64 split as 32+32) ---
pope_2d = AxialPoPE(dim=64, heads=8, axial_dims=(32, 32))

pos_emb_2d = pope_2d((16, 16))          # auto meshgrid → (256, 64) freqs
assert pos_emb_2d.freqs.shape == (256, 64)

q = torch.randn(1, 8, 256, 64)
k = torch.randn(1, 8, 256, 64)
rotated_q, rotated_k = AxialPoPE.apply_pope_to_qk(pos_emb_2d, q, k)
assert rotated_q.shape == (1, 8, 256, 128)
```
```

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

### apply_pope_to_qk Function

Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt

A standalone function (also a static method on `PoPE` and `AxialPoPE`) that applies a `PolarEmbedReturn` to query and key tensors. It supports partial rotation, complex-number output, and configurable magnitude activations.

```APIDOC
## apply_pope_to_qk — Standalone Function

### Description
`apply_pope_to_qk` is a standalone `@autocast`-safe function (also available as a static method on `PoPE` and `AxialPoPE`) that applies a `PolarEmbedReturn` to query and key tensors. It supports partial rotation (when `rotate_dim < head_dim`), complex-number output mode, and a configurable magnitude activation (default: `F.softplus`).

### Usage
```python
import torch
import torch.nn.functional as F
from PoPE_pytorch import PoPE
from PoPE_pytorch.pope import apply_pope_to_qk

pope = PoPE(dim=32, heads=4)  # only rotate the first 32 of 64 dims (partial)
pos_emb = pope(256)

q = torch.randn(2, 4, 256, 64)
k = torch.randn(2, 4, 256, 64)

# Partial rotation: first 32 dims rotated, remaining 32 passed through unchanged
rotated_q, rotated_k = apply_pope_to_qk(pos_emb, q, k)
assert rotated_q.shape == (2, 4, 256, 96)  # 32*2 + 32 = 96

# Using ReLU instead of softplus as magnitude activation
rotated_q2, rotated_k2 = apply_pope_to_qk(pos_emb, q, k, to_magnitude=F.relu)

# Return complex tensors instead of interleaved real
rotated_q_c, rotated_k_c = apply_pope_to_qk(pos_emb, q, k, return_complex=True)
assert rotated_q_c.is_complex()
```
```

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

### Compute Attention Similarity with PoPE

Source: https://github.com/hunterhogan/pope-pytorch/blob/main/README.md

Compute attention similarity using PoPE, avoiding expansion to avoid memory issues.

```python
import torch
from PoPE_pytorch import PoPE, compute_attn_similarity

# define pope

pope = PoPE(dim = 64, heads = 8).cuda()

# get rotations

pos_emb = pope(1024)

# queries and keys

q = torch.randn(1, 8, 1024, 64).cuda()
k = torch.randn(1, 8, 1024, 64).cuda()

# fused attention similarity, avoiding expanding 64 to 128

sim = compute_attn_similarity(q, k, pos_emb) # (1, 8, 1024, 1024)

attn = sim.softmax(dim = -1) # the usual in attention..
```

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

### PoPE Module

Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt

`PoPE` is a `torch.nn.Module` for computing polar-coordinate positional frequencies for 1D sequences. It can be instantiated and then called with a sequence length or position tensor to generate positional embeddings. It also provides a method to apply these embeddings to query and key tensors.

```APIDOC
## PoPE Module

### Description
`PoPE` is a `torch.nn.Module` that computes polar-coordinate positional frequencies for 1D sequences. It stores learned inverse frequencies and a per-head learnable phase bias clamped to `[-2π, 0]`. Calling `forward` with a sequence length (or explicit position tensor) returns a `PolarEmbedReturn` named tuple of `(freqs, bias)` ready to be applied to queries and keys.

### Instantiate
```python
import torch
from PoPE_pytorch import PoPE

# dim: per-head rotation dimension (must be <= head_dim in Q/K)
# heads: number of attention heads
pope = PoPE(dim=64, heads=8, theta=10000)
```

### Generate Positional Embeddings
```python
# Generate positional embeddings for a sequence of length 1024
pos_emb = pope(1024)
# pos_emb.freqs  shape: (1024, 64)
# pos_emb.bias   shape: (8, 64)
```

### Apply to Q and K (Training: Full Context)
```python
q = torch.randn(1, 8, 1024, 64)
k = torch.randn(1, 8, 1024, 64)

rotated_q, rotated_k = pope.apply_pope_to_qk(pos_emb, q, k)
assert rotated_q.shape == (1, 8, 1024, 128)  # dim doubles after polar interleaving
assert rotated_k.shape == (1, 8, 1024, 128)
```

### Apply to Q and K (Inference: Single New Token, Cached Keys)
```python
rotated_q_inf, rotated_k_inf = pope.apply_pope_to_qk(pos_emb, q[..., -1:, :], k)
assert rotated_q_inf.shape == (1, 8, 1,    128)
assert rotated_k_inf.shape == (1, 8, 1024, 128)
```

### Use Explicit Position Tensor
```python
# Use explicit position tensor (e.g. for packed/offset sequences)
positions = torch.arange(512, 512 + 1024).float()
pos_emb_offset = pope(positions)
```
```

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

### Compute Attention Similarity with PoPE

Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt

Computes the attention similarity matrix, optionally using a fused Triton kernel on CUDA for efficiency. Supports grouped query attention and different head layouts.

```python
import torch
from PoPE_pytorch import PoPE, compute_attn_similarity

pope = PoPE(dim=64, heads=8).cuda()
pos_emb = pope(1024)

q = torch.randn(1, 8, 1024, 64).cuda()
k = torch.randn(1, 8, 1024, 64).cuda()

# Fused similarity (uses Triton kernel on CUDA)
sim = compute_attn_similarity(q, k, pos_emb)         # (1, 8, 1024, 1024)
attn = sim.softmax(dim=-1)
```

```python
# Inference: single query token against full key cache
q_single = torch.randn(1, 8, 1, 64).cuda()
sim_single = compute_attn_similarity(q_single, k, pos_emb)   # (1, 8, 1, 1024)
```

```python
# Grouped Query Attention (4 KV heads, 8 Q heads)
pope_gqa = PoPE(dim=64, heads=4).cuda()
pos_emb_gqa = pope_gqa(1024)
q_gqa = torch.randn(1, 8, 1024, 64).cuda()
k_gqa = torch.randn(1, 4, 1024, 64).cuda()
sim_gqa = compute_attn_similarity(q_gqa, k_gqa, pos_emb_gqa)  # (1, 8, 1024, 1024)
```

```python
# Layout: heads-last  (batch, seq, heads, dim)
q_hl = torch.randn(1, 1024, 8, 64).cuda()
k_hl = torch.randn(1, 1024, 8, 64).cuda()
sim_hl = compute_attn_similarity(q_hl, k_hl, pos_emb, head_dimension_at_first=False)
```

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

### compute_attn_similarity

Source: https://context7.com/hunterhogan/pope-pytorch/llms.txt

Computes the attention similarity matrix after applying PoPE rotations. It can use a fused kernel on CUDA for efficiency or fall back to a standard einsum. Supports Grouped Query Attention (GQA).

```APIDOC
## compute_attn_similarity

### Description
Computes the full `(batch, heads, seq_q, seq_k)` attention similarity matrix after applying PoPE rotations. On CUDA with Triton available it uses a fused kernel that avoids expanding head_dim to 2×head_dim; it falls back to a standard einsum otherwise. Supports grouped query attention (GQA) when `q.heads > k.heads`.

### Method
`compute_attn_similarity(q, k, pos_emb, head_dimension_at_first=True)`

### Parameters
- **q** (torch.Tensor) - Query tensor.
- **k** (torch.Tensor) - Key tensor.
- **pos_emb** (torch.Tensor) - Positional embedding tensor.
- **head_dimension_at_first** (bool, optional) - Specifies if the head dimension is the first dimension after batch. Defaults to True.

### Request Example
```python
import torch
from PoPE_pytorch import PoPE, compute_attn_similarity

pope = PoPE(dim=64, heads=8).cuda()
pos_emb = pope(1024)

q = torch.randn(1, 8, 1024, 64).cuda()
k = torch.randn(1, 8, 1024, 64).cuda()

# Fused similarity (uses Triton kernel on CUDA)
sim = compute_attn_similarity(q, k, pos_emb)         # (1, 8, 1024, 1024)
attn = sim.softmax(dim=-1)

# Inference: single query token against full key cache
q_single = torch.randn(1, 8, 1, 64).cuda()
sim_single = compute_attn_similarity(q_single, k, pos_emb)   # (1, 8, 1, 1024)

# Grouped Query Attention (4 KV heads, 8 Q heads)
pope_gqa = PoPE(dim=64, heads=4).cuda()
pos_emb_gqa = pope_gqa(1024)
q_gqa = torch.randn(1, 8, 1024, 64).cuda()
k_gqa = torch.randn(1, 4, 1024, 64).cuda()
sim_gqa = compute_attn_similarity(q_gqa, k_gqa, pos_emb_gqa)  # (1, 8, 1024, 1024)

# Layout: heads-last  (batch, seq, heads, dim)
q_hl = torch.randn(1, 1024, 8, 64).cuda()
k_hl = torch.randn(1, 1024, 8, 64).cuda()
sim_hl = compute_attn_similarity(q_hl, k_hl, pos_emb, head_dimension_at_first=False)
```

### Response
#### Success Response (200)
- **sim** (torch.Tensor) - The computed attention similarity matrix.
```
Tokens:
Raw