faisalashraf/abaffinity

1# Copyright (c) Meta Platforms, Inc. and affiliates.2#3# This source code is licensed under the MIT license found in the4# LICENSE file in the root directory of this source tree.5 6from typing import Tuple7 8import torch9 10 11def rotate_half(x):12    x1, x2 = x.chunk(2, dim=-1)13    return torch.cat((-x2, x1), dim=-1)14 15 16def apply_rotary_pos_emb(x, cos, sin):17    cos = cos[:, : x.shape[-2], :]18    sin = sin[:, : x.shape[-2], :]19 20    return (x * cos) + (rotate_half(x) * sin)21 22 23class RotaryEmbedding(torch.nn.Module):24    """25    The rotary position embeddings from RoFormer_ (Su et. al).26    A crucial insight from the method is that the query and keys are27    transformed by rotation matrices which depend on the relative positions.28    Other implementations are available in the Rotary Transformer repo_ and in29    GPT-NeoX_, GPT-NeoX was an inspiration30    .. _RoFormer: https://arxiv.org/abs/2104.0986431    .. _repo: https://github.com/ZhuiyiTechnology/roformer32    .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox33    .. warning: Please note that this embedding is not registered on purpose, as it is transformative34        (it does not create the embedding dimension) and will likely be picked up (imported) on a ad-hoc basis35    """36 37    def __init__(self, dim: int, *_, **__):38        super().__init__()39        # Generate and save the inverse frequency buffer (non trainable)40        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))41        self.register_buffer("inv_freq", inv_freq)42 43        self._seq_len_cached = None44        self._cos_cached = None45        self._sin_cached = None46 47    def _update_cos_sin_tables(self, x, seq_dimension=1):48        seq_len = x.shape[seq_dimension]49 50        # Reset the tables if the sequence length has changed,51        # or if we're on a new device (possibly due to tracing for instance)52        if seq_len != self._seq_len_cached or self._cos_cached.device != x.device:53            self._seq_len_cached = seq_len54            t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(self.inv_freq)55            freqs = torch.einsum("i,j->ij", t, self.inv_freq)56            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)57 58            self._cos_cached = emb.cos()[None, :, :]59            self._sin_cached = emb.sin()[None, :, :]60 61        return self._cos_cached, self._sin_cached62 63    def forward(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:64        self._cos_cached, self._sin_cached = self._update_cos_sin_tables(k, seq_dimension=-2)65 66        return (67            apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached),68            apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached),69        )70