titans-pytorch 0.0.37__py3-none-any.whl → 0.0.39__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -7,7 +7,7 @@ from torch import nn, cat
 import torch.nn.functional as F
 from torch.nn import Module, ModuleList, Linear
 
-from einops import repeat, rearrange, pack, unpack
+from einops import einsum, repeat, rearrange, pack, unpack
 from einops.layers.torch import Rearrange
 
 from hyper_connections import get_init_and_expand_reduce_stream_functions
@@ -16,6 +16,7 @@ from hyper_connections import get_init_and_expand_reduce_stream_functions
 
 from axial_positional_embedding import ContinuousAxialPositionalEmbedding
 from rotary_embedding_torch import RotaryEmbedding
+from x_transformers.attend import Attend
 
 # proposed neural memory
 
@@ -93,6 +94,7 @@ class SegmentedAttention(Module):
         num_longterm_mem_tokens = 0,
         dim_head = 64,
         heads = 8,
+        attend_kwargs: dict = dict()
     ):
         super().__init__()
         self.norm = nn.RMSNorm(dim)
@@ -101,6 +103,8 @@ class SegmentedAttention(Module):
 
         self.rotary_emb = RotaryEmbedding(dim_head)
 
+        self.attend = Attend(causal = True, **attend_kwargs)
+
         self.to_qkv = LinearNoBias(dim, dim_inner * 3)
         self.to_out = LinearNoBias(dim_inner, dim)
 
@@ -145,9 +149,9 @@ class SegmentedAttention(Module):
         k = cat((pmk, k), dim = -2)
         v = cat((pmv, v), dim = -2)
 
-        # sdpa
+        # attention
 
-        out = F.scaled_dot_product_attention(q, k, v, is_causal = True)
+        out, _ = self.attend(q, k, v)
 
         out = self.merge_heads(out)
 
@@ -288,7 +292,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 +305,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/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(

{titans_pytorch-0.0.37.dist-info → titans_pytorch-0.0.39.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.37
+Version: 0.0.39
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -43,6 +43,7 @@ Requires-Dist: ninja
 Requires-Dist: rotary-embedding-torch
 Requires-Dist: tensordict
 Requires-Dist: torch>=2.2
+Requires-Dist: x-transformers
 Provides-Extra: examples
 Requires-Dist: local-attention>=1.10.1; extra == 'examples'
 Requires-Dist: taylor-series-linear-attention; extra == 'examples'

titans_pytorch-0.0.39.dist-info/RECORD

@@ -0,0 +1,9 @@
+titans_pytorch/__init__.py,sha256=F6pV8BamKCsbJFVo5x2hw69vzfJNLy54SwIKIueMdp4,142
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=h58sHfufxMnSXZXyWuW-KBwzq8xwBYmFjU2XtOjUixk,8512
+titans_pytorch/titans.py,sha256=bv2Ceq-_4nNb5FNx4hLd2inC93m5MmJxO2-Mbf6PKK4,14378
+titans_pytorch/titans_attn_memory.py,sha256=Rwx_-riGeISBefZg5Kjic8jzmmRRys-u93D2Kgb7Mos,12691
+titans_pytorch-0.0.39.dist-info/METADATA,sha256=3KD2hmJ-uOyQ87Z3VB6JfaKtDcLBnoKA8037DpzJuPE,3968
+titans_pytorch-0.0.39.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.39.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.39.dist-info/RECORD,,

titans_pytorch-0.0.37.dist-info/RECORD

@@ -1,9 +0,0 @@
-titans_pytorch/__init__.py,sha256=F6pV8BamKCsbJFVo5x2hw69vzfJNLy54SwIKIueMdp4,142
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=JjKGEMBit_SvyAsxq5v08614YBcLVx3OkM6pf0rADsA,8400
-titans_pytorch/titans.py,sha256=ALICGfc6p2bD2QkaibyIceVLvBIRKXmDm-w7RjnVOe4,14304
-titans_pytorch/titans_attn_memory.py,sha256=Rwx_-riGeISBefZg5Kjic8jzmmRRys-u93D2Kgb7Mos,12691
-titans_pytorch-0.0.37.dist-info/METADATA,sha256=RNokG8101_tlR0BiF-AxqYLZpXqafMSiN1Rg_pZe2-o,3938
-titans_pytorch-0.0.37.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.37.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.37.dist-info/RECORD,,