titans-pytorch 0.0.62__py3-none-any.whl → 0.0.64__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -153,6 +153,7 @@ class SegmentedAttention(Module):
         self.num_longterm_mem_tokens = num_longterm_mem_tokens
 
         total_segment_len = segment_len + num_longterm_mem_tokens
+        self.total_segment_len = total_segment_len
 
         self.split_heads = Rearrange('b n (h d) -> b h n d', h = heads)
         self.merge_heads = Rearrange('b h n d -> b n (h d)')
@@ -209,7 +210,7 @@ class SegmentedAttention(Module):
         # prep flex attention
 
         if not exists(flex_attn_fn):
-            block_mask = create_mac_block_mask(seq_len, self.segment_len, self.num_persist_mem_tokens)
+            block_mask = create_mac_block_mask(seq_len, self.total_segment_len, self.num_persist_mem_tokens)
 
             flex_attn_fn = partial(flex_attention, block_mask = block_mask)
 
@@ -241,7 +242,6 @@ class SegmentedAttention(Module):
         batch, seq_len = seq.shape[:2]
 
         # auto pad to multiple
-        # todo - get rid of logic with flex attention
 
         seq, inverse_segment = pad_and_segment_with_inverse(seq, total_segment_len)
 

titans_pytorch/titans.py

@@ -168,13 +168,14 @@ class MemoryAttention(Module):
     ):
         super().__init__()
         self.scale = scale
+        dim_ff_hidden = int(dim * expansion_factor)
 
         self.weights = nn.ParameterList([
             nn.Parameter(torch.randn(dim, dim)), # queries
             nn.Parameter(torch.randn(dim, dim)), # keys
             nn.Parameter(torch.randn(dim, dim)), # values
-            nn.Parameter(torch.randn(dim, dim * expansion_factor)), # ff w1
-            nn.Parameter(torch.randn(dim * expansion_factor, dim)), # ff w2
+            nn.Parameter(torch.randn(dim, dim_ff_hidden)), # ff w1
+            nn.Parameter(torch.randn(dim_ff_hidden, dim)), # ff w2
         ])
 
         for weight in self.weights:

{titans_pytorch-0.0.62.dist-info → titans_pytorch-0.0.64.dist-info}/METADATA

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

titans_pytorch-0.0.64.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=I-4oF1JPEmUvhLszEocM5cDgHYNFcNp0Q9nbDYSPFqU,195
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=7voYtbD_ErCD_JjvwhAiunUWtSIsIxGJAaf2aRB3c2s,15349
+titans_pytorch/titans.py,sha256=y6lJRErIoM6T2aTVFlf1GxSB0cpsmBZdSIj1DCHUCQ8,17486
+titans_pytorch-0.0.64.dist-info/METADATA,sha256=K63jobSfTdn-aFpEpZgolu4zSIvgUzF2rDuoCHGXkgE,4457
+titans_pytorch-0.0.64.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.64.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.64.dist-info/RECORD,,

titans_pytorch-0.0.62.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=I-4oF1JPEmUvhLszEocM5cDgHYNFcNp0Q9nbDYSPFqU,195
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=R9Xu-BjqAe9ZY60IGk4aNXBx_L8THsjJ4QrkbTnRNHo,15346
-titans_pytorch/titans.py,sha256=95J6UL44lOrdZSXdm7p36xC9tDeSmRBwdjig9T82PzI,17452
-titans_pytorch-0.0.62.dist-info/METADATA,sha256=08Blaa9Ehyv09rSA9uWguxbhKpbrd7C53Ya13E1VbpU,4457
-titans_pytorch-0.0.62.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.62.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.62.dist-info/RECORD,,