|
from typing import Optional, Tuple |
|
|
|
import torch |
|
|
|
|
|
def precompute_freqs_cis( |
|
dim: int, end: int, theta: float, device: Optional[torch.device] = None |
|
) -> torch.Tensor: |
|
freqs = 1.0 / ( |
|
theta ** (torch.arange(0, dim, 2, device=device)[: (dim // 2)].float() / dim) |
|
) |
|
t = torch.arange(end, device=freqs.device) |
|
freqs = torch.outer(t, freqs).float() |
|
return torch.polar(torch.ones_like(freqs), freqs) |
|
|
|
|
|
def apply_rotary_emb( |
|
xq: torch.Tensor, |
|
xk: torch.Tensor, |
|
freqs_cis: torch.Tensor, |
|
) -> Tuple[torch.Tensor, torch.Tensor]: |
|
xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) |
|
xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) |
|
freqs_cis = freqs_cis[:, None, :] |
|
xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(2) |
|
xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(2) |
|
return xq_out.type_as(xq), xk_out.type_as(xk) |
|
|