| |
|
|
| import torch |
| from typing import Tuple |
|
|
|
|
| def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0): |
| """ |
| Precompute the frequency tensor for complex exponentials (cis) with given dimensions. |
| |
| This function calculates a frequency tensor with complex exponentials using the given dimension 'dim' |
| and the end index 'end'. The 'theta' parameter scales the frequencies. |
| The returned tensor contains complex values in complex64 data type. |
| |
| Args: |
| dim (int): Dimension of the frequency tensor. |
| end (int): End index for precomputing frequencies. |
| theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0. |
| |
| Returns: |
| torch.Tensor: Precomputed frequency tensor with complex exponentials. |
| """ |
|
|
| freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (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 reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor): |
| assert freqs_cis.shape[1:] == (x.shape[1], x.shape[-1]) |
| return freqs_cis.contiguous().unsqueeze(2) |
|
|
|
|
| def apply_rotary_emb( |
| xq: torch.Tensor, |
| xk: torch.Tensor, |
| freqs_cis: torch.Tensor, |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| """ |
| Apply rotary embeddings to input tensors using the given frequency tensor. |
| |
| This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided |
| frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor |
| is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are |
| returned as real tensors. |
| |
| Args: |
| xq (torch.Tensor): Query tensor to apply rotary embeddings. |
| xk (torch.Tensor): Key tensor to apply rotary embeddings. |
| freqs_cis (torch.Tensor): Precomputed frequency tensor for complex exponentials. |
| |
| Returns: |
| Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings. |
| """ |
| 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 = reshape_for_broadcast(freqs_cis, xq_) |
| xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) |
| xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3) |
| return xq_out.type_as(xq), xk_out.type_as(xk) |
|
|
