titans-pytorch 0.3.22__tar.gz → 0.3.24__tar.gz

{titans_pytorch-0.3.22 → titans_pytorch-0.3.24}/PKG-INFO

@@ -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 → titans_pytorch-0.3.24}/pyproject.toml

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

{titans_pytorch-0.3.22 → titans_pytorch-0.3.24}/tests/test_titans.py

@@ -29,6 +29,7 @@ def torch_default_dtype(dtype):
 @pytest.mark.parametrize('chunk_size, attn_pool_chunks', ((64, True), (64, False), (1, False)))
 @pytest.mark.parametrize('momentum', (False, True))
 @pytest.mark.parametrize('qk_rmsnorm', (False, True))
+@pytest.mark.parametrize('heads', (1, 4))
 @pytest.mark.parametrize('max_grad_norm', (None, 2.))
 @pytest.mark.parametrize('num_kv_per_token', (1, 2))
 @pytest.mark.parametrize('per_parameter_lr_modulation', (False, True))
@@ -40,6 +41,7 @@ def test_titans(
     chunk_size,
     momentum,
     qk_rmsnorm,
+    heads,
     max_grad_norm,
     num_kv_per_token,
     per_parameter_lr_modulation,
@@ -54,6 +56,7 @@ def test_titans(
         num_kv_per_token = num_kv_per_token,
         momentum = momentum,
         qk_rmsnorm = qk_rmsnorm,
+        heads = heads,
         per_parameter_lr_modulation = per_parameter_lr_modulation,
         per_head_learned_parameters = per_head_learned_parameters
     )

{titans_pytorch-0.3.22 → titans_pytorch-0.3.24}/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