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