titans-pytorch 0.1.7__tar.gz → 0.1.8__tar.gz

{titans_pytorch-0.1.7 → titans_pytorch-0.1.8}/PKG-INFO

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

{titans_pytorch-0.1.7 → titans_pytorch-0.1.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.1.7"
+version = "0.1.8"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{titans_pytorch-0.1.7 → titans_pytorch-0.1.8}/tests/test_titans.py

@@ -53,6 +53,19 @@ def test_titans_attn_memory():
 
     assert seq.shape == retrieved.shape
 
+def test_retrieve_store_diff_seq():
+    mem = NeuralMemory(
+        dim = 384,
+        chunk_size = (64, 32),
+    )
+
+    retrieve_seq = torch.randn(2, 64 * 64, 384)
+    store_seq = torch.randn(2, 64 * 32, 384)
+
+    retrieved = mem(retrieve_seq, store_seq = store_seq)
+
+    assert retrieve_seq.shape == retrieved.shape
+
 @pytest.mark.parametrize('seq_len', (1023, 17))
 @pytest.mark.parametrize('num_persist_mem_tokens', (0, 16))
 @pytest.mark.parametrize('num_longterm_mem_tokens', (0, 16))

{titans_pytorch-0.1.7 → titans_pytorch-0.1.8}/titans_pytorch/titans.py

@@ -44,6 +44,9 @@ def default(v, d):
 def identity(t):
     return t
 
+def pair(v):
+    return (v, v) if not isinstance(v, tuple) else v
+
 def round_down_multiple(seq, mult):
     return seq // mult * mult
 
@@ -290,7 +293,7 @@ class NeuralMemory(Module):
     def __init__(
         self,
         dim,
-        chunk_size = 1,
+        chunk_size: int | tuple[int, int] = 1,
         dim_head = None,
         heads = 1,
         model: Module | None = None,
@@ -313,6 +316,8 @@ class NeuralMemory(Module):
         super().__init__()
         dim_head = default(dim_head, dim)
 
+        self.retrieve_chunk_size, self.store_chunk_size = pair(chunk_size)
+
         # norms
 
         self.retrieve_norm = nn.RMSNorm(dim) if pre_rmsnorm else nn.Identity()
@@ -380,6 +385,10 @@ class NeuralMemory(Module):
         self.empty_memory_embed = nn.Parameter(torch.zeros(dim))
         nn.init.normal_(self.empty_memory_embed, std = 0.02)
 
+        # `chunk_size` refers to chunk size used for storing to memory model weights
+
+        chunk_size = self.store_chunk_size
+
         # whether to use averaging of chunks, or attention pooling
 
         if not attn_pool_chunks:
@@ -451,11 +460,11 @@ class NeuralMemory(Module):
         past_state: tuple[dict[str, Tensor], dict[str, Tensor]],
         return_aux_kv_loss = False
     ):
-        seq_len = seq.shape[-2]
+        seq_len, chunk_size = seq.shape[-2], self.store_chunk_size
 
         # handle edge case
 
-        if seq_len < self.chunk_size:
+        if seq_len < chunk_size:
             past_weight, _ = past_state
             return TensorDict(past_weight).clone().zero_(), self.zero
 
@@ -464,8 +473,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, chunk_size = seq.shape[-2], self.chunk_size
-        round_down_seq_len = round_down_multiple(seq_len, self.chunk_size)
+        round_down_seq_len = round_down_multiple(seq_len, chunk_size)
 
         seq = seq[:, :round_down_seq_len]
 
@@ -597,12 +605,12 @@ class NeuralMemory(Module):
         seq,
         past_weights: dict[str, Tensor] | None = None,
     ):
-        chunk_size = self.chunk_size
+        chunk_size = self.retrieve_chunk_size
         batch, seq_len = seq.shape[:2]
 
         seq = self.retrieve_norm(seq)
 
-        if seq_len < self.chunk_size:
+        if seq_len < chunk_size:
             return self.init_empty_memory_embed(batch, seq_len)
 
         seq = seq[:, (chunk_size - 1):]
@@ -674,7 +682,7 @@ class NeuralMemory(Module):
     ):
         batch, seq_len = seq.shape[:2]
 
-        if seq_len < self.chunk_size:
+        if seq_len < self.retrieve_chunk_size:
             out = self.init_empty_memory_embed(batch, seq_len)
 
             if not return_aux_kv_loss: