titans-pytorch 0.0.35__py3-none-any.whl → 0.0.36__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -185,7 +185,9 @@ class MemoryAsContextTransformer(Module):
         # long term mem tokens
 
         self.segment_len = segment_len
+
         self.num_longterm_mem_tokens = num_longterm_mem_tokens
+        has_longterm_mems = num_longterm_mem_tokens > 0
 
         self.longterm_mems = nn.Parameter(torch.randn(num_longterm_mem_tokens, dim) * 0.02)
 
@@ -197,7 +199,11 @@ class MemoryAsContextTransformer(Module):
         self.neural_mem_layers = ModuleList([])
 
         layers = tuple(range(1, depth + 1))
-        neural_memory_layers = set(default(neural_memory_layers, layers))
+
+        if not exists(neural_memory_layers):
+            neural_memory_layers = layers if has_longterm_mems else ()
+
+        assert not (num_longterm_mem_tokens > 0 and len(neural_memory_layers) == 0), 'empty `neural_memory_layers` when longterm memory tokens are present'
 
         for layer in layers:
 
@@ -205,7 +211,9 @@ class MemoryAsContextTransformer(Module):
 
             mem = None
 
-            if num_longterm_mem_tokens > 0 and layer in neural_memory_layers:
+            if layer in neural_memory_layers:
+                assert has_longterm_mems, '`num_longterm_mem_tokens` must be greater than 0'
+
                 mem = NeuralMemory(dim = dim, chunk_size = num_longterm_mem_tokens)
                 mem = init_hyper_conn(dim = dim, branch = mem)
 

{titans_pytorch-0.0.35.dist-info → titans_pytorch-0.0.36.dist-info}/METADATA

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

{titans_pytorch-0.0.35.dist-info → titans_pytorch-0.0.36.dist-info}/RECORD

@@ -1,9 +1,9 @@
 titans_pytorch/__init__.py,sha256=F6pV8BamKCsbJFVo5x2hw69vzfJNLy54SwIKIueMdp4,142
 titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=FGShQHD-dQQdQKKzvNS_jTC_FcikdqO_s3ZKOKfr_9E,8502
+titans_pytorch/mac_transformer.py,sha256=xXQ9GvtvNArYidV1OOhUhCJ0pIxniElTLnL0_eIZtEE,8821
 titans_pytorch/titans.py,sha256=Kx_tl_QkeDccvkMwPZ0xQ_saYjZfbKzDNPTTSHNWYcA,14304
 titans_pytorch/titans_attn_memory.py,sha256=Rwx_-riGeISBefZg5Kjic8jzmmRRys-u93D2Kgb7Mos,12691
-titans_pytorch-0.0.35.dist-info/METADATA,sha256=jrhx-Bp1LqlOAV3jl4M70WTqq29ciz5lYWJvo2aoPE4,3938
-titans_pytorch-0.0.35.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.35.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.35.dist-info/RECORD,,
+titans_pytorch-0.0.36.dist-info/METADATA,sha256=7Bum0wO6e6BsB7TShLBZWALyStcs1LLaEv5vvnVlQ9c,3938
+titans_pytorch-0.0.36.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.36.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.36.dist-info/RECORD,,