titans-pytorch 0.3.11__py3-none-any.whl → 0.3.14__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -493,12 +493,20 @@ class MemoryAsContextTransformer(Module):
         use_flex_attn = False,
         sliding_window_attn = False,
         neural_mem_weight_residual = False,
+        token_emb: Module | None = None,
+        abs_pos_emb: Module | None = None
     ):
         super().__init__()
 
-        self.token_emb = nn.Embedding(num_tokens, dim)
+        if not exists(token_emb):
+            token_emb = nn.Embedding(num_tokens, dim)
 
-        self.axial_pos_emb = ContinuousAxialPositionalEmbedding(dim = dim, num_axial_dims = 2)
+        self.token_emb = token_emb
+
+        if not exists(abs_pos_emb):
+            abs_pos_emb = ContinuousAxialPositionalEmbedding(dim = dim, num_axial_dims = 2)
+
+        self.abs_pos_emb = abs_pos_emb
 
         # long term mem tokens
 

titans_pytorch/neural_memory.py

@@ -7,7 +7,7 @@ from itertools import zip_longest
 from collections import namedtuple
 
 import torch
-from torch import nn, cat, tensor, Tensor
+from torch import nn, stack, cat, tensor, Tensor
 import torch.nn.functional as F
 from torch.nn import Linear, Module, Parameter, ParameterList, ParameterDict
 from torch.func import functional_call, vmap, grad
@@ -230,6 +230,7 @@ class NeuralMemory(Module):
         momentum = True,
         momentum_order = 1,
         learned_momentum_combine = False,
+        learned_combine_include_zeroth = False,
         pre_rmsnorm = True,
         post_rmsnorm = False,
         qk_rmsnorm = False,
@@ -399,12 +400,17 @@ class NeuralMemory(Module):
             assert momentum
             assert momentum_order > 1, 'only second order momentum allowed for now, but may allow learned combination of zeroth'
 
+            if learned_combine_include_zeroth:
+                momentum_order += 1
+
             self.to_learned_momentum_combine = Sequential(
                 nn.Linear(dim, heads * momentum_order),
                 nn.Softmax(dim = -1),
                 Rearrange('b n (h o) -> o (b h) n', h = heads)
             )
 
+            self.learned_combine_include_zeroth = learned_combine_include_zeroth
+
         # per layer learning rate modulation
 
         self.to_layer_modulation = Sequential(
@@ -662,10 +668,14 @@ class NeuralMemory(Module):
 
                     momentums.append(momentum)
 
-                momentums = torch.stack(momentums)
+                momentums = stack(momentums)
 
                 next_last_momentum[param_name] = momentums[:, :, -1] # momentums shape is Float['o bh n 1']
 
+                if learned_combine and self.learned_combine_include_zeroth:
+                    # add the original surprise if learned combination of momentums
+                    momentums = cat((rearrange(surprise, '... -> 1 ...'), momentums), dim = 0)
+
                 if not learned_combine:
                     update = momentums[-1]
                 else:

{titans_pytorch-0.3.11.dist-info → titans_pytorch-0.3.14.dist-info}/METADATA

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

titans_pytorch-0.3.14.dist-info/RECORD

@@ -0,0 +1,9 @@
+titans_pytorch/__init__.py,sha256=AyEUlcXWpnqrvyeihRAXWIfQlzLA4NhBjOqQU4edL-4,297
+titans_pytorch/associative_scan.py,sha256=esaLbukFlgvy2aqopsqBy6KEcZ64B3rsNhG8moKdPSc,5159
+titans_pytorch/mac_transformer.py,sha256=F04B88GaH0wHseUIWaX6VFhOSsk_3XDQ1E8e6pvqKgQ,25170
+titans_pytorch/memory_models.py,sha256=wnH9i9kUSoVZhEWUlj8LpBSbB400L9kLt1zP8CO45QQ,5835
+titans_pytorch/neural_memory.py,sha256=VmUAS1xOM0ZfearWIzQrX_P7HI69viuwrg9M7BQByeE,29349
+titans_pytorch-0.3.14.dist-info/METADATA,sha256=1reoUZhvKaFPR6U0QXqJOziyss0HwHhwfJUf7oU7t-s,6817
+titans_pytorch-0.3.14.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.3.14.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.3.14.dist-info/RECORD,,

titans_pytorch-0.3.11.dist-info/RECORD

@@ -1,9 +0,0 @@
-titans_pytorch/__init__.py,sha256=AyEUlcXWpnqrvyeihRAXWIfQlzLA4NhBjOqQU4edL-4,297
-titans_pytorch/associative_scan.py,sha256=esaLbukFlgvy2aqopsqBy6KEcZ64B3rsNhG8moKdPSc,5159
-titans_pytorch/mac_transformer.py,sha256=EyqA53HBqvAr4UNZUs37LR6IltyEfA7FKEV54YzVYlg,24945
-titans_pytorch/memory_models.py,sha256=wnH9i9kUSoVZhEWUlj8LpBSbB400L9kLt1zP8CO45QQ,5835
-titans_pytorch/neural_memory.py,sha256=7YglrQaDpKS2hbpBBwx7PmqhJdjyvFEPZDt_QXmnUMM,28878
-titans_pytorch-0.3.11.dist-info/METADATA,sha256=xAEvavDiCj__5Bl_5UXaG__BycdUB2DzHOud-nwsn1c,6817
-titans_pytorch-0.3.11.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.3.11.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.3.11.dist-info/RECORD,,