titans-pytorch 0.0.4__tar.gz → 0.0.6__tar.gz

{titans_pytorch-0.0.4 → titans_pytorch-0.0.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.4
+Version: 0.0.6
 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
@@ -65,7 +66,10 @@ from titans_pytorch import NeuralMemory
 
 x = torch.randn(2, 64, 32)
 
-mem = NeuralMemory(32)
+mem = NeuralMemory(
+    dim = 32,
+    chunk_size = 2
+)
 
 out = mem(x)
 

{titans_pytorch-0.0.4 → titans_pytorch-0.0.6}/README.md

@@ -20,7 +20,10 @@ from titans_pytorch import NeuralMemory
 
 x = torch.randn(2, 64, 32)
 
-mem = NeuralMemory(32)
+mem = NeuralMemory(
+    dim = 32,
+    chunk_size = 2
+)
 
 out = mem(x)
 

{titans_pytorch-0.0.4 → titans_pytorch-0.0.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.0.4"
+version = "0.0.6"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -36,7 +36,9 @@ Homepage = "https://pypi.org/project/titans-pytorch/"
 Repository = "https://github.com/lucidrains/titans-pytorch"
 
 [project.optional-dependencies]
-examples = []
+examples = [
+    "local-attention>=1.9.15"
+]
 test = [
     "pytest"
 ]

{titans_pytorch-0.0.4 → titans_pytorch-0.0.6}/titans_pytorch/titans.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+import math
 from functools import partial
 
 import torch
@@ -11,7 +12,8 @@ from tensordict import TensorDict
 
 from titans_pytorch.associative_scan import (
     associative_scan,
-    binary_operator
+    binary_operator,
+    pad_at_dim
 )
 
 import einx
@@ -41,6 +43,9 @@ def default(v, d):
 def round_down_multiple(seq, mult):
     return seq // mult * mult
 
+def round_up_multiple(seq, mult):
+    return math.ceil(seq / mult) * mult
+
 def pack_one_with_inverse(t, pattern):
     packed, packed_shape = pack([t], pattern)
 
@@ -159,7 +164,7 @@ class NeuralMemory(Module):
         # 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]
@@ -229,14 +234,26 @@ 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,
         seq,
         past_weights: dict[str, Tensor] | None = None,
     ):
-        batch = seq.shape[0]
+        chunk_size = self.chunk_size
+        batch, seq_len = seq.shape[:2]
+
+        assert seq_len >= chunk_size
+
+        seq = seq[:, (chunk_size - 1):]
+        curtailed_seq_len = seq.shape[-2]
+
+        next_seq_len = round_up_multiple(curtailed_seq_len, chunk_size)
+
+        padding = next_seq_len - curtailed_seq_len
+
+        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
@@ -256,7 +273,7 @@ class NeuralMemory(Module):
         # fetch values from memory model
 
         curr_weights = curr_weights.apply(lambda t: rearrange(t, 'b n ... -> (b n) ...'))
-        queries = rearrange(queries, 'b n d -> (b n) 1 d')
+        queries = rearrange(queries, 'b (n c) d -> (b n) c d', c = chunk_size)
 
         # forward functional call
 
@@ -264,7 +281,12 @@ class NeuralMemory(Module):
 
         # reconstitute batch dimension
 
-        values = rearrange(values, '(b n) 1 d -> b n d', b = batch)
+        values = rearrange(values, '(b n) c d -> b (n c) d', b = batch)
+
+        # restore
+
+        values = pad_at_dim(values, (chunk_size - 1, 0), dim = 1, value = 0.) # todo, used a learned null memory embedding instead of 0s for retrieving from empty neural memory
+        values = values[:, :-padding]
 
         return values