rxnn 0.1.78__py3-none-any.whl → 0.1.80__py3-none-any.whl

rxnn/transformers/positional.py

@@ -12,14 +12,17 @@ class RotaryPositionalEmbedding(nn.Module):
         self.max_seq_len = max_seq_len
         self.base = base
         inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
-        self.register_buffer('inv_freq', inv_freq)
-        self.register_buffer('cache', None, persistent=False)
+        self.register_buffer('inv_freq', inv_freq) # must stay for models compatibility
+        # Pre-cache freqs for max_len
+        t = torch.arange(max_seq_len).type_as(self.inv_freq)
+        freqs = torch.einsum('i,j->ij', t, self.inv_freq)
+        self.register_buffer('cache', freqs)
+
 
     def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        device = q.device
         seq_len = q.size(-2)
         # Prepare RoPE Frequencies
-        freqs = self._prepare_freqs(seq_len, device)
+        freqs = self._prepare_freqs(seq_len)
 
         # Apply the rotation to the queries
         q_embed = self._rotate(q, freqs)
@@ -29,27 +32,17 @@ class RotaryPositionalEmbedding(nn.Module):
         return q_embed, k_embed
 
     def forward_one(self, q: torch.Tensor) -> torch.Tensor:
-        device = q.device
         seq_len = q.size(-2)
         # Prepare RoPE Frequencies
-        freqs = self._prepare_freqs(seq_len, device)
+        freqs = self._prepare_freqs(seq_len)
 
         # Apply the rotation to the queries
         q_embed = self._rotate(q, freqs)
 
         return q_embed
 
-    def _prepare_freqs(self, seq_len: int, device: torch.device) -> torch.Tensor:
-        cache_len = self.cache.size(1)
-        if self.cache is None or cache_len < seq_len:
-            t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
-            freqs = torch.einsum('i,j->ij', t, self.inv_freq)
-            self.cache = freqs
-            return freqs[None, None, :, :]
-        elif cache_len == seq_len:
-            return self.cache[None, None, :, :]
-        else:
-            return self.cache[:seq_len][None, None, :, :]
+    def _prepare_freqs(self, seq_len: int) -> torch.Tensor:
+        return self.cache[:seq_len][None, None, :, :]
 
     def _rotate(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
         x1 = x[..., 0::2]

{rxnn-0.1.78.dist-info → rxnn-0.1.80.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.1.78
+Version: 0.1.80
 Summary: RxNN: Reactive Neural Networks Platform
 License: Apache-2.0
 Keywords: deep-learning,ai,machine-learning

{rxnn-0.1.78.dist-info → rxnn-0.1.80.dist-info}/RECORD

@@ -22,10 +22,10 @@ rxnn/transformers/layers.py,sha256=OX8CsFY9A7uqH1SLwyexR_5BNlwheYrJHCGXjF8Q7HU,7
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
 rxnn/transformers/models.py,sha256=xbnn3FTNZFhaqq9A0XEM12ie_WL_58pPeq0qFXIgve0,7656
 rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
-rxnn/transformers/positional.py,sha256=DE1TP3D6ikBPg3Ym0sP9F666LHuE70H0w-JEH5DfKPw,4415
+rxnn/transformers/positional.py,sha256=ge-kaS6WnWnPGnWVp25ZK5bVkmhBUNCaELaN2rN_fSY,4097
 rxnn/transformers/sampler.py,sha256=poWBpxg1iuK5gEJtxHkk5VVfS9V48hs2Olqdhy_Gw8c,6548
 rxnn/utils.py,sha256=d5U8i5ukovgDyqiycc2AoxObTz_eF_bgo2MKvdtJ98s,467
-rxnn-0.1.78.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.78.dist-info/METADATA,sha256=559E3b22oEiu6vXNnsi7xLCw0GeuYQmcdmOgHkcdlL0,16589
-rxnn-0.1.78.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.78.dist-info/RECORD,,
+rxnn-0.1.80.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.80.dist-info/METADATA,sha256=Voy_a7EI9nC1kEqzxHcLCYpZRJntWnoKaFpF7XyiKCE,16589
+rxnn-0.1.80.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.80.dist-info/RECORD,,