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

titans_pytorch/mac_transformer.py

@@ -224,7 +224,7 @@ class MemoryAsContextTransformer(Module):
         self.layers = ModuleList([])
 
         self.neural_mem_layers = ModuleList([])
-        neural_memory_segment_len = default(neural_memory_segment_len, num_longterm_mem_tokens + segment_len)
+        self.neural_memory_segment_len = default(neural_memory_segment_len, num_longterm_mem_tokens + segment_len)
 
         layers = tuple(range(1, depth + 1))
 
@@ -245,7 +245,7 @@ class MemoryAsContextTransformer(Module):
 
                 mem = NeuralMemory(
                     dim = dim,
-                    chunk_size = neural_memory_segment_len,
+                    chunk_size = self.neural_memory_segment_len,
                     **neural_memory_kwargs
                 )
 
@@ -287,10 +287,7 @@ class MemoryAsContextTransformer(Module):
 
         # math
 
-        batch, seq_len, segment_len, num_longterm_mem_tokens= *x.shape, self.segment_len, self.num_longterm_mem_tokens
-
-        windows = ceil(seq_len / segment_len)
-        total_segment_len = segment_len + num_longterm_mem_tokens
+        batch, seq_len, neural_mem_segment_len, segment_len, num_longterm_mem_tokens = *x.shape, self.neural_memory_segment_len, self.segment_len, self.num_longterm_mem_tokens
 
         # token embedding
 
@@ -305,11 +302,16 @@ class MemoryAsContextTransformer(Module):
 
         x = inverse_segment(x)
 
+        seq_len_with_mem = x.shape[-2]
+
         # apply axial positional embedding
         # so intra and inter segment can be more easily discerned by the network
 
-        pos_emb = self.axial_pos_emb((windows, total_segment_len), flatten = True)
-        x = x + pos_emb[:x.shape[-2]]
+        neural_mem_windows = ceil(seq_len_with_mem / neural_mem_segment_len)
+
+        pos_emb = self.axial_pos_emb((neural_mem_windows, neural_mem_segment_len), flatten = True)
+
+        x = x + pos_emb[:seq_len_with_mem]
 
         # value residual
 
@@ -334,7 +336,7 @@ class MemoryAsContextTransformer(Module):
 
         # excise out the memories
 
-        x, inverse_segment = pad_and_segment_with_inverse(x, total_segment_len)
+        x, inverse_segment = pad_and_segment_with_inverse(x, segment_len + num_longterm_mem_tokens)
 
         x, _ = unpack(x, mem_ps, 'b * d')
 

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

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

titans_pytorch-0.0.44.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=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,,

titans_pytorch-0.0.43.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=kSdfWGWwEk6d0lbb0WLVKQwdmG8LAzDg36QZm7aIio0,9451
-titans_pytorch/titans.py,sha256=qxQ8pZCz8GEDhKeJMEaeAEzH66GAGVBNaRdNam_-czg,15260
-titans_pytorch-0.0.43.dist-info/METADATA,sha256=3Rlt_5CIeDUkYEK5tcLiWTseWv48gg4OH5vMoSVNS2w,4210
-titans_pytorch-0.0.43.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.43.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.43.dist-info/RECORD,,