titans-pytorch 0.2.28__py3-none-any.whl → 0.3.1__py3-none-any.whl

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))
 
@@ -658,13 +652,14 @@ class NeuralMemory(Module):
         next_last_update = TensorDict()
         next_last_momentum = TensorDict()
 
-        for (param_name, surprise), (_, last_update), (_, last_momentum) in zip(surprises.items(), past_last_update.items(), past_last_momentum.items()):
+        for (param_name, surprise), (_, last_update) in zip(surprises.items(), past_last_update.items()):
 
             update = surprise
 
             # derive momentum with associative scan - eq (10)
 
             if has_momentum:
+                last_momentum = past_last_momentum[param_name]
                 update = self.assoc_scan(adaptive_momentum, surprise, prev = last_momentum) # momentum is S / surprise in the paper
                 momentum = update
                 next_last_momentum[param_name] = momentum[:, -1]
@@ -872,15 +867,20 @@ class NeuralMemory(Module):
             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
 
-            past_state = (weights, last_momentum)
-            next_neural_mem_state = tuple_index_set(next_neural_mem_state, -2, past_state)
-            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
 
@@ -891,7 +891,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(

{titans_pytorch-0.2.28.dist-info → titans_pytorch-0.3.1.dist-info}/METADATA

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

titans_pytorch-0.3.1.dist-info/RECORD

@@ -0,0 +1,9 @@
+titans_pytorch/__init__.py,sha256=Y3m_ZlpEqYwp-Md1ARhNGJxq8bQp8ty1o039nZOOJo0,276
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=5rO4GQxSyFWWEc3pc3xNyG0sK5EXE7MmxKI-_kEMl2M,24941
+titans_pytorch/memory_models.py,sha256=Q9SAIyAbStF5Tz0EhvRbn3yAdE3nk3xKc1ndieIe714,4671
+titans_pytorch/neural_memory.py,sha256=K1z7wtv366Y6-eEyXMFZ_j7D2frPl5RxfSgxzFYoFMc,27704
+titans_pytorch-0.3.1.dist-info/METADATA,sha256=ZAxucKq2DZBtW-BI_O2sUQ5RXy11a7eu48yPpwnanpw,6815
+titans_pytorch-0.3.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.3.1.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.3.1.dist-info/RECORD,,

titans_pytorch-0.2.28.dist-info/RECORD

@@ -1,9 +0,0 @@
-titans_pytorch/__init__.py,sha256=Y3m_ZlpEqYwp-Md1ARhNGJxq8bQp8ty1o039nZOOJo0,276
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=5rO4GQxSyFWWEc3pc3xNyG0sK5EXE7MmxKI-_kEMl2M,24941
-titans_pytorch/memory_models.py,sha256=Q9SAIyAbStF5Tz0EhvRbn3yAdE3nk3xKc1ndieIe714,4671
-titans_pytorch/neural_memory.py,sha256=vFkyxO2WK-mEnf-KP3Za4oJ7_c1KrLHonCuelofJtFs,27754
-titans_pytorch-0.2.28.dist-info/METADATA,sha256=HxSE1da0QnLMbJLpsO2-2wJYm_kFUmdezy2JCUEx6aU,6816
-titans_pytorch-0.2.28.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.2.28.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.2.28.dist-info/RECORD,,