digital-garbageandClaude Opus 5.5
53cbc29948
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
32 lines
1.4 KiB
Python
32 lines
1.4 KiB
Python
"""H3's rotary position encoding, undone and redone outside the model.
|
|
|
|
H3 turns the first `2*half` dims of every head's query by a per-row 2x2 matrix
|
|
(split-half: dim i pairs with i+half). A feature that averages queries across
|
|
rows, or adds one vector to all of them, has to work where the rows are NOT yet
|
|
turned -- otherwise the average cancels and the push points somewhere different
|
|
at every position. Mirrors comfy_kitchen's eager `apply_rope_split_half1`.
|
|
"""
|
|
|
|
import torch
|
|
|
|
|
|
def query_rotation(rope_freqs, seq_len):
|
|
"""-> rotate(x [..., n, D], rows slice, inverse=False), or None when
|
|
`rope_freqs` is not H3's [1, S, 1, half, 2, 2] table for this sequence."""
|
|
if not torch.is_tensor(rope_freqs) or rope_freqs.ndim != 6 \
|
|
or rope_freqs.shape[1] != seq_len or tuple(rope_freqs.shape[-2:]) != (2, 2):
|
|
return None
|
|
table = rope_freqs[0, :, 0] # [S, half, 2, 2]
|
|
|
|
def rotate(x, rows, inverse=False):
|
|
t = table[rows]
|
|
half = int(t.shape[-3])
|
|
rot = 2 * half
|
|
pairs = x[..., :rot].reshape(*x.shape[:-1], 2, half).movedim(-2, -1).float()
|
|
mat = t.float().transpose(-1, -2) if inverse else t.float()
|
|
out = (mat @ pairs.unsqueeze(-1)).squeeze(-1)
|
|
out = out.movedim(-1, -2).reshape(*x.shape[:-1], rot).to(x.dtype)
|
|
return torch.cat([out, x[..., rot:]], dim=-1)
|
|
|
|
return rotate
|