titans-pytorch 0.1.29__tar.gz → 0.1.31__tar.gz

{titans_pytorch-0.1.29 → titans_pytorch-0.1.31}/PKG-INFO

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

{titans_pytorch-0.1.29 → titans_pytorch-0.1.31}/pyproject.toml

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

{titans_pytorch-0.1.29 → titans_pytorch-0.1.31}/tests/test_titans.py

@@ -124,18 +124,22 @@ def test_mac(
     assert logits.shape == (1, seq_len, 256)
 
 @pytest.mark.parametrize('sliding', (False, True))
-@pytest.mark.parametrize('mem_layers', ((), None, (4,)))
+@pytest.mark.parametrize('mem_layers', (()))
+@pytest.mark.parametrize('longterm_mems', (0, 4, 16))
+@pytest.mark.parametrize('prompt_len', (0, 4, 16))
 def test_mac_sampling(
     sliding,
-    mem_layers
+    mem_layers,
+    longterm_mems,
+    prompt_len
 ):
     transformer = MemoryAsContextTransformer(
         num_tokens = 256,
         dim = 256,
-        depth = 2,
+        depth = 4,
         segment_len = 32,
         num_persist_mem_tokens = 4,
-        num_longterm_mem_tokens = 0,
+        num_longterm_mem_tokens = longterm_mems,
         sliding_window_attn = sliding,
         neural_memory_layers = mem_layers,
         neural_mem_gate_attn_output = False
@@ -145,8 +149,10 @@ def test_mac_sampling(
 
     # after much training
 
-    sampled = transformer.sample(ids[:, :4], 53, use_cache = False, temperature = 0.)
-    sampled_with_cache = transformer.sample(ids[:, :4], 53, use_cache = True, temperature = 0.)
+    prompt = ids[:, :prompt_len]
+
+    sampled = transformer.sample(prompt, 53, use_cache = False, temperature = 0.)
+    sampled_with_cache = transformer.sample(prompt, 53, use_cache = True, temperature = 0.)
 
     assert torch.allclose(sampled, sampled_with_cache)
 

{titans_pytorch-0.1.29 → titans_pytorch-0.1.31}/titans_pytorch/mac_transformer.py

@@ -582,13 +582,8 @@ class MemoryAsContextTransformer(Module):
         self,
         seq_index
     ):
-        total_segment_len = self.attn_window_size
-
-        seq = seq_index + 1
-        seq -= int((seq % total_segment_len) == 0)
-        last_segment_len = round_down_multiple(seq, total_segment_len)
-        segment_seq = seq - last_segment_len
-        return (segment_seq - self.segment_len) > 0
+        total_segment_len, segment_len = self.attn_window_size, self.segment_len
+        return ((seq_index % total_segment_len + 1) - segment_len) > 0
 
     def seq_len_with_longterm_mem(
         self,
@@ -597,7 +592,7 @@ class MemoryAsContextTransformer(Module):
         assert seq_len > 0
 
         segment_len, num_mem = self.segment_len, self.num_longterm_mem_tokens
-        return ceil(seq_len / segment_len) * num_mem + seq_len
+        return ((seq_len - 1) // segment_len) * num_mem + seq_len
 
     @torch.no_grad()
     def sample(
@@ -723,9 +718,9 @@ class MemoryAsContextTransformer(Module):
         is_inferencing = exists(cache)
 
         if not exists(cache):
-            cache = (None, None)
+            cache = (seq_len_with_mem - 1, None, None)
 
-        kv_caches, neural_mem_caches = cache
+        inference_seq_index, kv_caches, neural_mem_caches = cache
 
         kv_caches = iter(default(kv_caches, []))
         neural_mem_caches = iter(default(neural_mem_caches, []))
@@ -744,7 +739,8 @@ class MemoryAsContextTransformer(Module):
         # when inferencing, only do one token at a time
 
         if is_inferencing:
-            x = x[:, -1:]
+            ind = inference_seq_index
+            x = x[:, ind:(ind + 1)]
 
         # expand and reduce streams for hyper connections
 
@@ -817,6 +813,17 @@ class MemoryAsContextTransformer(Module):
             if not self.sliding_window_attn and divisible_by(seq_len_with_mem, attn_window_size):
                 next_kv_caches = next_kv_caches[..., 0:0, :]
 
+            next_cache = (
+                inference_seq_index + 1,
+                next_kv_caches,
+                next_neural_mem_caches
+            )
+
+            is_longterm_mem = self.seq_index_is_longterm(inference_seq_index)
+
+            if is_inferencing and is_longterm_mem:
+                return None, next_cache
+
         # hyper connection reducing of streams
 
         x = self.reduce_streams(x)
@@ -829,7 +836,7 @@ class MemoryAsContextTransformer(Module):
 
             x, _ = inverse_pack_mems(x)
 
-            x = inverse_segment(x)
+            x = inverse_segment(x, remove_pad = False)
 
             x = x[:, :seq_len]
 
@@ -843,7 +850,7 @@ class MemoryAsContextTransformer(Module):
             if not return_cache:
                 return logits
 
-            return logits, (next_kv_caches, next_neural_mem_caches)
+            return logits, next_cache
 
         ar_loss = F.cross_entropy(rearrange(logits, 'b n l -> b l n'), labels)
 

{titans_pytorch-0.1.29 → titans_pytorch-0.1.31}/titans_pytorch/titans.py

@@ -409,6 +409,7 @@ class NeuralMemory(Module):
     ):
         super().__init__()
         dim_head = default(dim_head, dim)
+        assert not (heads == 1 and dim_head != dim)
 
         self.retrieve_chunk_size, self.store_chunk_size = pair(chunk_size)
 
@@ -566,7 +567,7 @@ class NeuralMemory(Module):
     ):
         assert xnor(exists(value_residual), exists(self.learned_value_residual))
 
-        seq_len, chunk_size = seq.shape[-2], default(chunk_size, self.store_chunk_size)
+        seq_len, heads, chunk_size = seq.shape[-2], self.heads, default(chunk_size, self.store_chunk_size)
 
         # handle edge case
 
@@ -645,7 +646,7 @@ class NeuralMemory(Module):
 
         # restore batch and sequence dimension
 
-        grads = grads.apply(lambda t: rearrange(t, '(b n) ... -> b n ...', b = batch))
+        grads = grads.apply(lambda t: rearrange(t, '(b n) ... -> b n ...', b = batch * heads))
 
         # maybe per layer modulation