titans-pytorch 0.2.18__py3-none-any.whl → 0.2.19__py3-none-any.whl

titans_pytorch/neural_memory.py

@@ -66,6 +66,12 @@ 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))
 
@@ -296,7 +302,7 @@ class NeuralMemory(Module):
         init_adaptive_step_bias = None,
         init_momentum_bias = None,
         init_decay_bias = None,
-        learned_weight_residual = False,
+        accept_weight_residual = False,
         default_model_kwargs: dict = dict(
             depth = 2
         )
@@ -445,7 +451,7 @@ class NeuralMemory(Module):
             nn.Linear(dim, heads),
             Rearrange('b n h -> b h n'),
             nn.Sigmoid()
-        ) if learned_weight_residual else None
+        ) if accept_weight_residual else None
 
         # allow for softclamp the gradient norms for storing memories
 
@@ -584,7 +590,6 @@ class NeuralMemory(Module):
             if exists(self.to_learned_weight_residual_mix):
                 mix = self.to_learned_weight_residual_mix(chunked_seq)
                 mix = rearrange(mix, 'b h n -> (b h) n')
-
                 prev_weights = prev_weights.apply(lambda t: einx.multiply('bh n, bh n ... -> bh n ...', mix, t))
 
             weights_for_surprise = weights_for_surprise + prev_weights
@@ -854,9 +859,9 @@ class NeuralMemory(Module):
 
             weights = last_update
 
-            next_neural_mem_state = list(next_neural_mem_state)
-            next_neural_mem_state[1] = last_update
-            next_neural_mem_state = NeuralMemCache(*next_neural_mem_state)
+            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, updates)
 
         # retrieve
 

{titans_pytorch-0.2.18.dist-info → titans_pytorch-0.2.19.dist-info}/METADATA

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

{titans_pytorch-0.2.18.dist-info → titans_pytorch-0.2.19.dist-info}/RECORD

@@ -2,8 +2,8 @@ titans_pytorch/__init__.py,sha256=Y3m_ZlpEqYwp-Md1ARhNGJxq8bQp8ty1o039nZOOJo0,27
 titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
 titans_pytorch/mac_transformer.py,sha256=J9MHFViWWehBOFOKx7ry_X2k8nAXaAZFFeCwGtudZyk,24942
 titans_pytorch/memory_models.py,sha256=Q9SAIyAbStF5Tz0EhvRbn3yAdE3nk3xKc1ndieIe714,4671
-titans_pytorch/neural_memory.py,sha256=A74v2g4OUXJIqXLu1Rk7FDsTlhveJIHsiamsWN9e7cI,26602
-titans_pytorch-0.2.18.dist-info/METADATA,sha256=QFJI4OoZZSdgZd6FJInAG5-z2ubpvO5I0aCRNKwhDnQ,6812
-titans_pytorch-0.2.18.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.2.18.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.2.18.dist-info/RECORD,,
+titans_pytorch/neural_memory.py,sha256=e1TENsK0IptpapFRtQux-Uii2MFvisoiQLzTQboHc50,26709
+titans_pytorch-0.2.19.dist-info/METADATA,sha256=6frwjwq4CX7VcdCMh8Op9WeDHLRD-oeo8DlTde538UI,6812
+titans_pytorch-0.2.19.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.2.19.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.2.19.dist-info/RECORD,,