titans-pytorch 0.3.22__py3-none-any.whl → 0.3.24__py3-none-any.whl

titans_pytorch/neural_memory.py

@@ -289,6 +289,8 @@ class NeuralMemory(Module):
         self.heads = heads
 
         self.split_heads = Rearrange('b n (h d) -> b h n d', h = heads)
+        self.split_kv_heads = Rearrange('b n (h u d) -> b h (n u) d', h = heads, u = num_kv_per_token)
+
         self.merge_heads = Rearrange('b h n d -> b n (h d)')
         self.combine_heads = LinearNoBias(dim_inner, dim) if heads > 1 else nn.Identity()
 
@@ -425,8 +427,8 @@ class NeuralMemory(Module):
 
             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)
+                Rearrange('b n (h o) -> o (b h) n', h = heads),
+                nn.Softmax(dim = 0),
             )
 
             self.learned_combine_include_zeroth = learned_combine_include_zeroth
@@ -596,17 +598,15 @@ class NeuralMemory(Module):
 
         # maybe multi head
 
-        keys, values = map(self.split_heads, (keys, values))
-
-        batch = keys.shape[0]
+        keys, values = map(self.split_kv_heads, (keys, values))
 
-        # take care of chunking
+        # maybe keys rmsnorm
 
-        keys, values = tuple(rearrange(t, 'b h (n c) (u d) -> (b h n) (c u) d', c = chunk_size, u = num_updates) for t in (keys, values))
+        keys = self.k_norm(keys)
 
-        # maybe qk rmsnorm
+        # take care of chunking
 
-        keys = self.k_norm(keys)
+        keys, values = tuple(rearrange(t, 'b h (n c u) d -> (b h n) (c u) d', c = chunk_size, u = num_updates) for t in (keys, values))
 
         # adaptive lr
 

{titans_pytorch-0.3.22.dist-info → titans_pytorch-0.3.24.dist-info}/METADATA

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

{titans_pytorch-0.3.22.dist-info → titans_pytorch-0.3.24.dist-info}/RECORD

@@ -2,8 +2,8 @@ titans_pytorch/__init__.py,sha256=AyEUlcXWpnqrvyeihRAXWIfQlzLA4NhBjOqQU4edL-4,29
 titans_pytorch/associative_scan.py,sha256=esaLbukFlgvy2aqopsqBy6KEcZ64B3rsNhG8moKdPSc,5159
 titans_pytorch/mac_transformer.py,sha256=grD327B3OCIy7d23jNUWIoUo1bIgXUqD26dXWCjdi28,25565
 titans_pytorch/memory_models.py,sha256=wnH9i9kUSoVZhEWUlj8LpBSbB400L9kLt1zP8CO45QQ,5835
-titans_pytorch/neural_memory.py,sha256=2NPaesGpSrYoGnp4VshYPm1gvA1W7xcuJzZ2nydCcGY,31161
-titans_pytorch-0.3.22.dist-info/METADATA,sha256=kOPNtBdpt8QQMENoPhjE1wqN-zQLrt40QCqtMTatRhQ,6817
-titans_pytorch-0.3.22.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.3.22.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.3.22.dist-info/RECORD,,
+titans_pytorch/neural_memory.py,sha256=haSepQdfGsQdoo9Yk5agvRR91kTu8kgkXpBmBZaH8WI,31237
+titans_pytorch-0.3.24.dist-info/METADATA,sha256=0-WHTKNXZpESWfOMSOO8MiWddqDoSRP1lifsfgHmewo,6817
+titans_pytorch-0.3.24.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.3.24.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.3.24.dist-info/RECORD,,