titans-pytorch 0.2.25__tar.gz → 0.2.27__tar.gz

{titans_pytorch-0.2.25 → titans_pytorch-0.2.27}/PKG-INFO

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

{titans_pytorch-0.2.25 → titans_pytorch-0.2.27}/pyproject.toml

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

{titans_pytorch-0.2.25 → titans_pytorch-0.2.27}/tests/test_titans.py

@@ -73,12 +73,16 @@ def test_titans_attn_memory():
 
     assert seq.shape == retrieved.shape
 
-def test_neural_mem_chaining_chunks():
+@pytest.mark.parametrize('gated_transition', (True, False))
+def test_neural_mem_chaining_chunks(
+    gated_transition
+):
     mem  = NeuralMemory(
         dim = 384,
         dim_head = 64,
         heads = 2,
-        chunk_size = 16
+        chunk_size = 16,
+        gated_transition = gated_transition
     )
 
     seq = torch.randn(2, 48, 384)

{titans_pytorch-0.2.25 → titans_pytorch-0.2.27}/titans_pytorch/neural_memory.py

@@ -872,11 +872,11 @@ class NeuralMemory(Module):
             last_update, _ = past_state
 
             if exists(gate):
-                weights = TensorDict({param_name: v1.lerp(v2, gate) for (param_name, v1), (_, v2) in zip(weights.items(), last_update.items())})
+                weights = TensorDict({param_name: one_weight.lerp(one_last_update, gate) for (param_name, one_weight), (_, one_last_update) in zip(weights.items(), last_update.items())})
             else:
                 weights = last_update
 
-            next_neural_mem_state = tuple_index_set(next_neural_mem_state, 1, last_update)
+            next_neural_mem_state = tuple_index_set(next_neural_mem_state, 1, weights)
 
         next_neural_mem_state = tuple_index_set(next_neural_mem_state, -1, updates)