| 152 | |
| 153 | |
| 154 | class RotaryPositionalEmbedding1D(nn.Module): |
| 155 | |
| 156 | def __init__(self, |
| 157 | head_dim, |
| 158 | ): |
| 159 | super().__init__() |
| 160 | self.head_dim = head_dim |
| 161 | self.base = 10000 |
| 162 | |
| 163 | |
| 164 | @lru_cache(maxsize=32) |
| 165 | def precompute_freqs_cis_1d(self, pos_indices): |
| 166 | |
| 167 | freqs = 1.0 / (self.base ** (torch.arange(0, self.head_dim, 2)[: (self.head_dim // 2)].float() / self.head_dim)) |
| 168 | freqs = freqs.to(pos_indices.device) |
| 169 | freqs = torch.einsum("..., f -> ... f", pos_indices.float(), freqs) |
| 170 | freqs = repeat(freqs, "... n -> ... (n r)", r=2) |
| 171 | return freqs |
| 172 | |
| 173 | def forward(self, x, pos_indices): |
| 174 | """1D RoPE. |
| 175 | |
| 176 | Args: |
| 177 | query (torch.tensor): [B, head, seq, head_dim] |
| 178 | pos_indices (torch.tensor): [seq,] |
| 179 | Returns: |
| 180 | query with the same shape as input. |
| 181 | """ |
| 182 | freqs_cis = self.precompute_freqs_cis_1d(pos_indices) |
| 183 | |
| 184 | x_ = x.float() |
| 185 | |
| 186 | freqs_cis = freqs_cis.float().to(x.device) |
| 187 | cos, sin = freqs_cis.cos(), freqs_cis.sin() |
| 188 | cos, sin = rearrange(cos, 'n d -> 1 1 n d'), rearrange(sin, 'n d -> 1 1 n d') |
| 189 | x_ = (x_ * cos) + (rotate_half(x_) * sin) |
| 190 | |
| 191 | return x_.type_as(x) |
| 192 | |
| 193 | |
| 194 | |