"""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)