| 226 | |
| 227 | |
| 228 | class SinusoidsPositionEmbedding(nn.Module): |
| 229 | def __init__(self, length, channels, max_timescale=10000): |
| 230 | super().__init__() |
| 231 | if channels % 2 != 0: |
| 232 | raise ValueError("SinusoidsPositionEmbedding needs even channels input") |
| 233 | log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1) |
| 234 | inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2).float()) |
| 235 | scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :] |
| 236 | self.register_buffer( |
| 237 | "positional_embedding", |
| 238 | torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1), |
| 239 | persistent=False, |
| 240 | ) |
| 241 | |
| 242 | def forward(self, seqlen: int): |
| 243 | return self.positional_embedding[:seqlen, :] |
| 244 | |
| 245 | |
| 246 | @auto_docstring( |