(dim: int, end: int, theta: float = 10000.0)
| 36 | |
| 37 | |
| 38 | def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0): |
| 39 | freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) |
| 40 | t = torch.arange(end, device=freqs.device) # type: ignore |
| 41 | freqs = torch.outer(t, freqs).float() # type: ignore |
| 42 | freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 |
| 43 | return freqs_cis |
| 44 | |
| 45 | |
| 46 | def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor): |