titans-pytorch 0.0.37__tar.gz → 0.0.38__tar.gz

{titans_pytorch-0.0.37 → titans_pytorch-0.0.38}/PKG-INFO

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

{titans_pytorch-0.0.37 → titans_pytorch-0.0.38}/pyproject.toml

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

{titans_pytorch-0.0.37 → titans_pytorch-0.0.38}/titans_pytorch/mac_transformer.py

@@ -288,7 +288,8 @@ class MemoryAsContextTransformer(Module):
         for (attn, ff), maybe_neural_mem in zip(self.layers, self.neural_mem_layers):
 
             if exists(maybe_neural_mem):
-                mems = maybe_neural_mem(mems)
+                x = maybe_neural_mem(x)
+
 
             x = attn(x)
 
@@ -300,7 +301,7 @@ class MemoryAsContextTransformer(Module):
 
         x, inverse_segment = pad_and_segment_with_inverse(x, total_segment_len)
 
-        x, mem = unpack(x, mem_ps, 'b * d')
+        x, _ = unpack(x, mem_ps, 'b * d')
 
         x = inverse_segment(x)
 

{titans_pytorch-0.0.37 → titans_pytorch-0.0.38}/titans_pytorch/titans.py

@@ -27,9 +27,7 @@ n - sequence
 d - feature dimension
 c - intra-chunk
 """
-
-# constants
-
+7
 LinearNoBias = partial(Linear, bias = False)
 
 # functions
@@ -390,7 +388,10 @@ class NeuralMemory(Module):
 
         padding = next_seq_len - curtailed_seq_len
 
-        seq = pad_at_dim(seq, (0, padding), dim = 1)
+        needs_pad = padding > 0
+
+        if needs_pad:
+            seq = pad_at_dim(seq, (0, padding), dim = 1)
 
         # the parameters of the memory model stores the memories of the key / values
         # when the MLP has only 1 weight matrix, it is equivalent to `kv` fast weight memories from linear attention literature (recall fetching of memories is q @ (kv)) / schmidhuber's paper
@@ -442,7 +443,9 @@ class NeuralMemory(Module):
         empty_memory_embeds = self.init_empty_memory_embed(values.shape[0], chunk_size - 1)
         values = torch.cat((empty_memory_embeds, values), dim = -2)
 
-        values = values[:, :-padding]
+        if needs_pad:
+            values = values[:, :-padding]
+
         return values
 
     def forward(