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

titans_pytorch/__init__.py

@@ -3,4 +3,6 @@ from titans_pytorch.titans import (
     MemoryMLP,
 )
 
-from titans_pytorch.mac_transformer import MemoryAsContextTransformer
+from titans_pytorch.mac_transformer import (
+    MemoryAsContextTransformer
+)

titans_pytorch/titans.py

@@ -425,11 +425,7 @@ class NeuralMemory(Module):
         next_seq_len = round_up_multiple(curtailed_seq_len, chunk_size)
 
         padding = next_seq_len - curtailed_seq_len
-
-        needs_pad = padding > 0
-
-        if needs_pad:
-            seq = pad_at_dim(seq, (0, padding), dim = 1)
+        seq = pad_at_dim(seq, (0, padding), dim = 1)
 
         # the parameters of the memory model stores the memories of the key / values
         # when the MLP has only 1 weight matrix, it is equivalent to `kv` fast weight memories from linear attention literature (recall fetching of memories is q @ (kv)) / schmidhuber's paper
@@ -481,10 +477,7 @@ class NeuralMemory(Module):
         empty_memory_embeds = self.init_empty_memory_embed(values.shape[0], chunk_size - 1)
         values = torch.cat((empty_memory_embeds, values), dim = -2)
 
-        if needs_pad:
-            values = values[:, :-padding]
-
-        return values
+        return values[:, :seq_len]
 
     def forward(
         self,

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.42
+Version: 0.0.43
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -83,6 +83,26 @@ retrieved = mem(seq)
 assert seq.shape == retrieved.shape
 ```
 
+A transformer with the `MAC` configuration can be used as
+
+```python
+import torch
+from titans_pytorch import MemoryAsContextTransformer
+
+transformer = MemoryAsContextTransformer(
+    num_tokens = 256,
+    dim = 256,
+    depth = 2,
+    segment_len = 128,              # local attention window size
+    num_persist_mem_tokens = 4,
+    num_longterm_mem_tokens = 16,
+)
+
+token_ids = torch.randint(0, 256, (1, 1023))
+
+logits = transformer(token_ids) # (1, 1023, 256)
+```
+
 ## Experiments
 
 ```bash

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

titans_pytorch-0.0.42.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=F6pV8BamKCsbJFVo5x2hw69vzfJNLy54SwIKIueMdp4,142
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=kSdfWGWwEk6d0lbb0WLVKQwdmG8LAzDg36QZm7aIio0,9451
-titans_pytorch/titans.py,sha256=eA7D9aqfGbtmC2SgGAQnfEVYp5Uza9uebEyDpVpjNQc,15372
-titans_pytorch-0.0.42.dist-info/METADATA,sha256=4lZBFMZPuQRDQGdTK-TWheytxEaQZQv7bdXO7MLBrwI,3744
-titans_pytorch-0.0.42.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.42.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.42.dist-info/RECORD,,