titans-pytorch 0.1.8__tar.gz → 0.1.10__tar.gz

{titans_pytorch-0.1.8 → titans_pytorch-0.1.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.1.8
+Version: 0.1.10
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -35,7 +35,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Requires-Dist: accelerated-scan>=0.2.0
-Requires-Dist: axial-positional-embedding>=0.3.5
+Requires-Dist: axial-positional-embedding>=0.3.6
 Requires-Dist: einops>=0.8.0
 Requires-Dist: einx>=0.3.0
 Requires-Dist: hyper-connections>=0.1.8
@@ -62,7 +62,7 @@ Unofficial implementation of [Titans](https://arxiv.org/abs/2501.00663) in Pytor
 
 ## Appreciation
 
-- [@sentialx](https://github.com/sentialx) for sharing his early experimental results with me
+- [Eryk](https://github.com/sentialx) for sharing his early experimental results with me, positive for 2 layer MLP
 
 ## Install
 

{titans_pytorch-0.1.8 → titans_pytorch-0.1.10}/README.md

@@ -8,7 +8,7 @@ Unofficial implementation of [Titans](https://arxiv.org/abs/2501.00663) in Pytor
 
 ## Appreciation
 
-- [@sentialx](https://github.com/sentialx) for sharing his early experimental results with me
+- [Eryk](https://github.com/sentialx) for sharing his early experimental results with me, positive for 2 layer MLP
 
 ## Install
 

{titans_pytorch-0.1.8 → titans_pytorch-0.1.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.1.8"
+version = "0.1.10"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -26,7 +26,7 @@ classifiers=[
 
 dependencies = [
     "accelerated-scan>=0.2.0",
-    "axial_positional_embedding>=0.3.5",
+    "axial_positional_embedding>=0.3.6",
     "einops>=0.8.0",
     "einx>=0.3.0",
     "hyper-connections>=0.1.8",

{titans_pytorch-0.1.8 → titans_pytorch-0.1.10}/titans_pytorch/mac_transformer.py

@@ -83,6 +83,14 @@ def identity(t):
 def round_up_multiple(seq, mult):
     return ceil(seq / mult) * mult
 
+def pack_with_inverse(t, pattern):
+    packed, packed_shape = pack(t, pattern)
+
+    def inverse(out, inv_pattern = None):
+        return unpack(out, packed_shape, default(inv_pattern, pattern))
+
+    return packed, inverse
+
 def pad_at_dim(t, pad, dim = -1, value = 0.):
     dims_from_right = (- dim - 1) if dim < 0 else (t.ndim - dim - 1)
     zeros = ((0, 0) * dims_from_right)
@@ -576,7 +584,7 @@ class MemoryAsContextTransformer(Module):
         x, inverse_segment = pad_and_segment_with_inverse(x, segment_len)
 
         mems = repeat(self.longterm_mems, 'n d -> b n d', b = x.shape[0])
-        x, mem_ps = pack((x, mems), 'b * d')
+        x, inverse_pack_mems = pack_with_inverse((x, mems), 'b * d')
 
         x = inverse_segment(x)
 
@@ -585,11 +593,9 @@ class MemoryAsContextTransformer(Module):
         # apply axial positional embedding
         # so intra and inter segment can be more easily discerned by the network
 
-        neural_mem_windows = ceil(seq_len_with_mem / neural_mem_segment_len)
-
-        pos_emb = self.axial_pos_emb((neural_mem_windows, neural_mem_segment_len), flatten = True)
+        pos_emb = self.axial_pos_emb.forward_with_seq_len(seq_len_with_mem, (neural_mem_segment_len,))
 
-        x = x + pos_emb[:seq_len_with_mem]
+        x = x + pos_emb
 
         # prep flex attention
 
@@ -634,7 +640,7 @@ class MemoryAsContextTransformer(Module):
 
         x, inverse_segment = pad_and_segment_with_inverse(x, segment_len + num_longterm_mem_tokens)
 
-        x, _ = unpack(x, mem_ps, 'b * d')
+        x, _ = inverse_pack_mems(x)
 
         x = inverse_segment(x)
 

{titans_pytorch-0.1.8 → titans_pytorch-0.1.10}/titans_pytorch/titans.py

@@ -391,6 +391,8 @@ class NeuralMemory(Module):
 
         # whether to use averaging of chunks, or attention pooling
 
+        assert not (attn_pool_chunks and chunk_size == 1), '`attn_pool_chunks` cannot be set to True if `chunk_size` is set to 1'
+
         if not attn_pool_chunks:
             chunk_reduce_module = Reduce('b (n c) ... -> b n ...', 'mean', c = chunk_size)
         else: