titans-pytorch 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

titans_pytorch/titans.py

@@ -91,10 +91,14 @@ class NeuralMemory(Module):
         dim,
         chunk_size = 1,
         model: Module | None = None,
-        store_memory_loss_fn: Callable = default_loss_fn
+        store_memory_loss_fn: Callable = default_loss_fn,
+        pre_rmsnorm = False
     ):
         super().__init__()
 
+        self.retrieve_norm = nn.RMSNorm(dim) if pre_rmsnorm else nn.Identity()
+        self.store_norm = nn.RMSNorm(dim) if pre_rmsnorm else nn.Identity()
+
         if not exists(model):
             model = MLP(dim, depth = 4)
 
@@ -161,10 +165,12 @@ class NeuralMemory(Module):
         past_state: tuple[dict[str, Tensor], dict[str, Tensor]]
     ):
 
+        seq = self.store_norm(seq)
+
         # curtail sequence by multiple of the chunk size
         # only a complete chunk of the sequence provides the memory for the next chunk
 
-        seq_len = seq.shape[-2]
+        seq_len, chunk_size = seq.shape[-2], self.chunk_size
         round_down_seq_len = round_down_multiple(seq_len, self.chunk_size)
 
         seq = seq[:, :round_down_seq_len]
@@ -234,7 +240,7 @@ class NeuralMemory(Module):
 
         next_state = (curr_weights + last_update, next_momentum)
 
-        return updates, next_state, aux_store_loss.mean()
+        return updates, next_state, aux_store_loss.mean() / chunk_size
 
     def retrieve_memories(
         self,
@@ -244,6 +250,8 @@ class NeuralMemory(Module):
         chunk_size = self.chunk_size
         batch, seq_len = seq.shape[:2]
 
+        seq = self.retrieve_norm(seq)
+
         assert seq_len >= chunk_size
 
         seq = seq[:, (chunk_size - 1):]

{titans_pytorch-0.0.5.dist-info → titans_pytorch-0.0.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.5
+Version: 0.0.7
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -39,6 +39,7 @@ Requires-Dist: einx>=0.3.0
 Requires-Dist: tensordict>=0.6.2
 Requires-Dist: torch>=2.3
 Provides-Extra: examples
+Requires-Dist: local-attention>=1.9.15; extra == 'examples'
 Provides-Extra: test
 Requires-Dist: pytest; extra == 'test'
 Description-Content-Type: text/markdown

titans_pytorch-0.0.7.dist-info/RECORD

@@ -0,0 +1,7 @@
+titans_pytorch/__init__.py,sha256=QKuJPCOJCdgtaPeKoHEkYkiQe65_LV9_8-cIMbBPU30,55
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/titans.py,sha256=IuAhvpeEGP_XQfhpWGEwNcmlSUKmtmtxWjjgwdEy0oI,9730
+titans_pytorch-0.0.7.dist-info/METADATA,sha256=NAIbruJrJLaafP5SjcM3a_6qe6yCgiohpC09CzOKsMg,3092
+titans_pytorch-0.0.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.7.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.7.dist-info/RECORD,,

titans_pytorch-0.0.5.dist-info/RECORD

@@ -1,7 +0,0 @@
-titans_pytorch/__init__.py,sha256=QKuJPCOJCdgtaPeKoHEkYkiQe65_LV9_8-cIMbBPU30,55
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/titans.py,sha256=3Mewuysj0g7iAlfjdqMlJhn9-pKJuOerB1frQmQYXuc,9428
-titans_pytorch-0.0.5.dist-info/METADATA,sha256=f1DgCKZz9nqNfZOrqbOpyn-yEx2v5M5zgGIW0Zeu84I,3032
-titans_pytorch-0.0.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.5.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.5.dist-info/RECORD,,