(self, pos)
| 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) |