| 246 | |
| 247 | |
| 248 | class AxialRoPE(nn.Module): |
| 249 | def __init__(self, dim, n_heads): |
| 250 | super().__init__() |
| 251 | log_min = math.log(math.pi) |
| 252 | log_max = math.log(10.0 * math.pi) |
| 253 | freqs = torch.linspace(log_min, log_max, n_heads * dim // 4 + 1)[:-1].exp() |
| 254 | self.register_buffer("freqs", freqs.view(dim // 4, n_heads).T.contiguous()) |
| 255 | |
| 256 | def extra_repr(self): |
| 257 | return f"dim={self.freqs.shape[1] * 4}, n_heads={self.freqs.shape[0]}" |
| 258 | |
| 259 | def forward(self, pos): |
| 260 | theta_h = pos[..., None, 0:1] * self.freqs.to(pos.dtype) |
| 261 | theta_w = pos[..., None, 1:2] * self.freqs.to(pos.dtype) |
| 262 | return torch.cat((theta_h, theta_w), dim=-1) |
| 263 | |
| 264 | |
| 265 | # Shifted window attention |