titans-pytorch 0.0.30__py3-none-any.whl → 0.0.31__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -29,9 +29,33 @@ def exists(v):
 def default(v, d):
     return v if exists(v) else d
 
+def identity(t):
+    return t
+
 def round_up_multiple(seq, mult):
     return ceil(seq / mult) * mult
 
+def pad_and_segment_with_inverse(seq, segment_len):
+    batch, seq_len = seq.shape[:2]
+
+    need_segment = seq_len >= segment_len
+
+    if not need_segment:
+        return seq, identity
+
+    next_seq_len_mult = round_up_multiple(seq_len, segment_len)
+
+    padding = next_seq_len_mult - seq_len
+    seq = F.pad(seq, (0, 0, 0, padding))
+
+    seq = rearrange(seq, 'b (w n) d -> (b w) n d', n = segment_len)
+
+    def inverse(out):
+        out = rearrange(out, '(b w) n d -> b (w n) d', b = batch)
+        return out[:, :-padding]
+
+    return seq, inverse
+
 # feedforward and attention
 
 class GEGLU(Module):
@@ -77,9 +101,6 @@ class SegmentedAttention(Module):
         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)')
 
-        self.segment_seq = Rearrange('b (n w) d -> (b n) w d', n = total_segment_len)
-        self.merge_seq_back = Rearrange('(b n) w d -> b (n w) d', n = total_segment_len)
-
         self.persistent_memory = nn.Parameter(torch.zeros(2, heads, num_persist_mem_tokens, dim_head))
 
     def forward(self, seq):
@@ -91,16 +112,7 @@ class SegmentedAttention(Module):
         # auto pad to multiple
         # todo - get rid of logic with flex attention
 
-        need_segment = seq_len >= total_segment_len
-
-        if need_segment:
-            next_seq_len = round_up_multiple(seq_len, total_segment_len)
-            padding = next_seq_len - seq_len
-
-            if padding > 0:
-                seq = F.pad(seq, (0, 0, 0, padding))
-
-            seq = self.segment_seq(seq)
+        seq, inverse_segment = pad_and_segment_with_inverse(seq, total_segment_len)
 
         # attention
 
@@ -130,10 +142,9 @@ class SegmentedAttention(Module):
 
         out = self.to_out(out)
 
-        if need_segment:
-            out = self.merge_seq_back(out)
+        out = inverse_segment(out)
 
-        return out[:, :seq_len]
+        return out
 
 # MAC transformer
 
@@ -207,29 +218,18 @@ class MemoryAsContextTransformer(Module):
 
         # intersperse longterm memory
 
-        need_segment = seq_len >= segment_len
-
-        if need_segment:
-            next_seq_len = round_up_multiple(seq_len, segment_len)
-            padding = next_seq_len - seq_len
-
-            if padding > 0:
-                x = F.pad(x, (0, 0, 0, padding))
-
-            x = rearrange(x, 'b (w n) d -> (b w) n d', n = segment_len)
+        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 = torch.cat((mems, x), dim = -2)
 
-        if need_segment:
-            x = rearrange(x, '(b w) n d -> b (w n) d', b = batch)
-            x = x[:, :seq_len]
+        x = inverse_segment(x)
 
         # apply axial positional embedding
         # so intra and inter segment can be more easily discerned by the network
 
         pos_emb = self.axial_pos_emb((windows, total_segment_len), flatten = True)
-        x = x + pos_emb[:seq_len]
+        x = x + pos_emb[:x.shape[-2]]
 
         # expand and reduce streams for hyper connections
 
@@ -241,6 +241,14 @@ class MemoryAsContextTransformer(Module):
 
         x = self.reduce_streams(x)
 
+        # excise out the memories
+
+        x, inverse_segment = pad_and_segment_with_inverse(x, total_segment_len)
+
+        x = x[:, self.num_longterm_mem_tokens:]
+
+        x = inverse_segment(x)
+
         # to logits
 
         x = self.norm(x)

{titans_pytorch-0.0.30.dist-info → titans_pytorch-0.0.31.dist-info}/METADATA

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

{titans_pytorch-0.0.30.dist-info → titans_pytorch-0.0.31.dist-info}/RECORD

@@ -1,9 +1,9 @@
 titans_pytorch/__init__.py,sha256=nB0873FZ_OyCda3qFeWTdpO4LbbrXDEvtAefVLzh6g0,71
 titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=XDLc9NWXpVMza03XjU0lkw5lRvtJ25ReTPKNoGslOOk,6773
+titans_pytorch/mac_transformer.py,sha256=bq5RbCgA0GWLFHTrDTIKUSQhkkuCkdjEykOwjfHDs0M,6747
 titans_pytorch/titans.py,sha256=Kx_tl_QkeDccvkMwPZ0xQ_saYjZfbKzDNPTTSHNWYcA,14304
 titans_pytorch/titans_attn_memory.py,sha256=Rwx_-riGeISBefZg5Kjic8jzmmRRys-u93D2Kgb7Mos,12691
-titans_pytorch-0.0.30.dist-info/METADATA,sha256=7H6WPsgfBE9ByEUP7r6C-cfBX6K13yNCaVhZMzwUvf8,3938
-titans_pytorch-0.0.30.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.30.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.30.dist-info/RECORD,,
+titans_pytorch-0.0.31.dist-info/METADATA,sha256=bN1fVL2S_vML1oqLIA92tvBhkVvnpQN11fU4e1QVI4s,3938
+titans_pytorch-0.0.31.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.31.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.31.dist-info/RECORD,,