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

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

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

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

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

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

@@ -220,16 +220,19 @@ def test_mac_sampling(
 @pytest.mark.parametrize('seq_len', (2, 64, 256))
 @pytest.mark.parametrize('prompt_len', (0, 65))
 @pytest.mark.parametrize('mem_chunk_size', (2, 32, 64))
+@pytest.mark.parametrize('gated_transition', (False, True))
 @torch_default_dtype(torch.float64)
 def test_neural_mem_inference(
     seq_len,
     prompt_len,
-    mem_chunk_size
+    mem_chunk_size,
+    gated_transition
 ):
 
     mem = NeuralMemory(
         dim = 384,
         chunk_size = mem_chunk_size,
+        gated_transition = gated_transition
     )
 
     seq = torch.randn(2, seq_len, 384)

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

@@ -66,12 +66,6 @@ def xnor(x, y):
 def divisible_by(num, den):
     return (num % den) == 0
 
-def tuple_index_set(t: tuple, index, value):
-    klass = type(t)
-    t = list(t)
-    t[index] = value
-    return klass(*t)
-
 def safe_cat(inputs, dim = -2):
     inputs = tuple(filter(exists, inputs))
 
@@ -869,16 +863,23 @@ class NeuralMemory(Module):
 
             # update weights once batch size is fulfilled
 
-            last_update, _ = past_state
+            last_update, last_momentum = past_state
 
             if exists(gate):
-                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
+                last_update = 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())})
+
+            past_state = (last_update, last_momentum)
+
+            # set weights to the last updated weights for the last minibatch
 
-            next_neural_mem_state = tuple_index_set(next_neural_mem_state, 1, weights)
+            weights = last_update
 
-        next_neural_mem_state = tuple_index_set(next_neural_mem_state, -1, updates)
+            next_neural_mem_state = next_neural_mem_state._replace(
+                weights = weights,
+                states = past_state,
+            )
+
+        next_neural_mem_state = next_neural_mem_state._replace(updates = updates)
 
         # retrieve
 
@@ -889,7 +890,8 @@ class NeuralMemory(Module):
             retrieve_chunk_size = 1
             need_pad = False
 
-            last_update, _ = past_state
+            last_update, _ = next_neural_mem_state.states
+
             updates = rearrange_dict_values(last_update, 'b ... -> b 1 ...')
 
         retrieved = self.retrieve_memories(