| 84 | |
| 85 | |
| 86 | class AxialRoPE(nn.Module): |
| 87 | def __init__(self, dim, n_heads, start_index=0, freqs_init=freqs_pixel_log(max_freq=10.0)): |
| 88 | super().__init__() |
| 89 | self.n_heads = n_heads |
| 90 | self.start_index = start_index |
| 91 | log_freqs = freqs_init((n_heads, dim // 4)) |
| 92 | self.freqs_h = nn.Parameter(log_freqs.clone()) |
| 93 | self.freqs_w = nn.Parameter(log_freqs.clone()) |
| 94 | |
| 95 | def extra_repr(self): |
| 96 | dim = (self.freqs_h.shape[-1] + self.freqs_w.shape[-1]) * 2 |
| 97 | return f"dim={dim}, n_heads={self.n_heads}, start_index={self.start_index}" |
| 98 | |
| 99 | def get_freqs(self, pos): |
| 100 | if pos.shape[-1] != 2: |
| 101 | raise ValueError("input shape must be (..., 2)") |
| 102 | freqs_h = pos[..., None, None, 0] * self.freqs_h.exp() |
| 103 | freqs_w = pos[..., None, None, 1] * self.freqs_w.exp() |
| 104 | freqs = torch.cat((freqs_h, freqs_w), dim=-1).repeat_interleave(2, dim=-1) |
| 105 | return freqs.transpose(-2, -3) |
| 106 | |
| 107 | def forward(self, x, pos): |
| 108 | freqs = self.get_freqs(pos) |
| 109 | return apply_rotary_emb(freqs, x, self.start_index) |
nothing calls this directly
no outgoing calls
no test coverage detected