titans-pytorch 0.0.44__py3-none-any.whl → 0.0.45__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -311,7 +311,7 @@ class MemoryAsContextTransformer(Module):
 
         pos_emb = self.axial_pos_emb((neural_mem_windows, neural_mem_segment_len), flatten = True)
 
-        x = x + pos_emb[:seq_len_with_mem]
+        pos_emb = pos_emb[:seq_len_with_mem]
 
         # value residual
 
@@ -324,7 +324,7 @@ class MemoryAsContextTransformer(Module):
         for (attn, ff), maybe_neural_mem in zip(self.layers, self.neural_mem_layers):
 
             if exists(maybe_neural_mem):
-                x = maybe_neural_mem(x)
+                x = maybe_neural_mem(x, pos_emb = pos_emb)
 
             x, values = attn(x, value_residual = value_residual)
 

titans_pytorch/titans.py

@@ -484,10 +484,14 @@ class NeuralMemory(Module):
         seq,
         store_seq = None,
         past_state: tuple[dict[str, Tensor], dict[str, Tensor]] | None = None,
-        return_next_memories = False
+        return_next_memories = False,
+        pos_emb: Tensor | None = None
     ):
         batch, seq_len = seq.shape[:2]
 
+        if exists(pos_emb):
+            seq = seq + pos_emb
+
         if seq_len < self.chunk_size:
             return self.init_empty_memory_embed(batch, seq_len)
 

{titans_pytorch-0.0.44.dist-info → titans_pytorch-0.0.45.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.44
+Version: 0.0.45
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch

titans_pytorch-0.0.45.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=wnv_Cjdjqh_h5IqLkQ8xrTtA2K663ITEn-1JeeHofTo,150
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=SFB7sXDt1bYpwt_PVrXM0-1vXKEemBTAfnfboU66A7M,9586
+titans_pytorch/titans.py,sha256=7LZIbaavC0bk85UBPzNzZP6YxKeFb0ujZ9k5IU048aI,15360
+titans_pytorch-0.0.45.dist-info/METADATA,sha256=EqrDXchEvzFbz1BqSdAB8HkPMjRY3KYyBSu16hbKTUs,4210
+titans_pytorch-0.0.45.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.45.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.45.dist-info/RECORD,,

titans_pytorch-0.0.44.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=wnv_Cjdjqh_h5IqLkQ8xrTtA2K663ITEn-1JeeHofTo,150
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=mF8PYAjeAjLas1gkYybgzZX1AVK82A_ps_LY00ofYYs,9565
-titans_pytorch/titans.py,sha256=qxQ8pZCz8GEDhKeJMEaeAEzH66GAGVBNaRdNam_-czg,15260
-titans_pytorch-0.0.44.dist-info/METADATA,sha256=QzgJ6YqBfRMvCMUnMPCtiil0Q488hL9BAfHhMEmI2pA,4210
-titans_pytorch-0.0.44.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.44.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.44.dist-info/RECORD,,