16 lines
717 B
Python
16 lines
717 B
Python
|
"""H3 three-axis split-half rotary position embeddings."""
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
|
||
|
|
def h3_rope_rotation(position_ids: torch.Tensor, inv_freq: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
|
||
|
|
"""Build H3's `[1, sequence, 1, 48, 2, 2]` rotation table."""
|
||
|
|
positions = position_ids.to(torch.float32)
|
||
|
|
frequencies = inv_freq.to(device=positions.device, dtype=torch.float32)
|
||
|
|
per_axis = positions.unsqueeze(-1) * frequencies.view(1, 1, -1)
|
||
|
|
half_angles = torch.cat(per_axis.unbind(dim=1), dim=-1)
|
||
|
|
cosine, sine = torch.cos(half_angles), torch.sin(half_angles)
|
||
|
|
return torch.stack((cosine, -sine, sine, cosine), dim=-1).reshape(
|
||
|
|
1, positions.shape[0], 1, half_angles.shape[-1], 2, 2
|
||
|
|
).to(dtype)
|