1D RoPE. Args: query (torch.tensor): [B, head, seq, head_dim] pos_indices (torch.tensor): [seq,] Returns: query with the same shape as input.
(self, x, pos_indices)
| 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 |
nothing calls this directly
no test coverage detected