(max_seq_len, dim, theta=10000)
| 30 | |
| 31 | @amp.autocast(enabled=False) |
| 32 | def rope_params(max_seq_len, dim, theta=10000): |
| 33 | assert dim % 2 == 0 |
| 34 | freqs = torch.outer( |
| 35 | torch.arange(max_seq_len), |
| 36 | 1.0 / torch.pow(theta, |
| 37 | torch.arange(0, dim, 2).to(torch.float64).div(dim))) |
| 38 | freqs = torch.polar(torch.ones_like(freqs), freqs) |
| 39 | return freqs |
| 40 | |
| 41 | |
| 42 | @amp.autocast(enabled=False) |