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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.1.7
+Version: 0.1.9
 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.9}/pyproject.toml

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

{titans_pytorch-0.1.7 → titans_pytorch-0.1.9}/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.9}/titans_pytorch/mac_transformer.py

@@ -83,6 +83,14 @@ def identity(t):
 def round_up_multiple(seq, mult):
     return ceil(seq / mult) * mult
 
+def pack_with_inverse(t, pattern):
+    packed, packed_shape = pack(t, pattern)
+
+    def inverse(out, inv_pattern = None):
+        return unpack(out, packed_shape, default(inv_pattern, pattern))
+
+    return packed, inverse
+
 def pad_at_dim(t, pad, dim = -1, value = 0.):
     dims_from_right = (- dim - 1) if dim < 0 else (t.ndim - dim - 1)
     zeros = ((0, 0) * dims_from_right)
@@ -576,7 +584,7 @@ class MemoryAsContextTransformer(Module):
         x, inverse_segment = pad_and_segment_with_inverse(x, segment_len)
 
         mems = repeat(self.longterm_mems, 'n d -> b n d', b = x.shape[0])
-        x, mem_ps = pack((x, mems), 'b * d')
+        x, inverse_pack_mems = pack_with_inverse((x, mems), 'b * d')
 
         x = inverse_segment(x)
 
@@ -634,7 +642,7 @@ class MemoryAsContextTransformer(Module):
 
         x, inverse_segment = pad_and_segment_with_inverse(x, segment_len + num_longterm_mem_tokens)
 
-        x, _ = unpack(x, mem_ps, 'b * d')
+        x, _ = inverse_pack_mems(x)
 
         x = inverse_segment(x)
 

{titans_pytorch-0.1.7 → titans_pytorch-0.1.9}/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,8 +385,14 @@ 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
 
+        assert not (attn_pool_chunks and chunk_size == 1), '`attn_pool_chunks` cannot be set to True if `chunk_size` is set to 1'
+
         if not attn_pool_chunks:
             chunk_reduce_module = Reduce('b (n c) ... -> b n ...', 'mean', c = chunk_size)
         else:
@@ -451,11 +462,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 +475,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 +607,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 +684,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: