import torch
import torch.nn as nn

class RoPE(nn.Module):
    def __init__(self, d_head: int, max_seq_len: int = 32768):
        super().__init__()
        assert d_head % 2 == 0

        inv_freq = 1.0 / (10000 ** (torch.arange(0, d_head, 2, dtype=torch.float32) / d_head))
        t = torch.arange(max_seq_len, dtype=torch.float32)
        freqs = torch.outer(t, inv_freq) 

        cos = freqs.cos().unsqueeze(0).unsqueeze(0)
        sin = freqs.sin().unsqueeze(0).unsqueeze(0)

        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: (B, H, S, D)
        """
        seq_len = x.shape[2]
        cos = self.cos[..., :seq_len, :].to(x.dtype)
        sin = self.sin[..., :seq_len, :].to(x.dtype)

        x1, x2 = x.chunk(2, dim=-1)                
        return torch.cat((x1 * cos - x2 * sin,
                          x1 * sin + x2 * cos), dim=-1)