titans-pytorch 0.0.32__py3-none-any.whl → 0.0.34__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -50,13 +50,20 @@ def pad_and_segment_with_inverse(seq, segment_len):
     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))
+    needs_pad = padding > 0
+
+    if needs_pad:
+        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]
+
+        if needs_pad:
+            out = out[:, :-padding]
+
+        return out
 
     return seq, inverse
 
@@ -226,7 +233,14 @@ class MemoryAsContextTransformer(Module):
 
         self.to_logits = LinearNoBias(dim, num_tokens)
 
-    def forward(self, x):
+    def forward(
+        self,
+        x,
+        return_loss = False
+    ):
+
+        if return_loss:
+            x, labels = x[:, :-1], x[:, 1:]
 
         # math
 
@@ -262,6 +276,7 @@ class MemoryAsContextTransformer(Module):
 
             if exists(maybe_neural_mem):
                 batch_streams = x.shape[0]
+
                 x, inverse_segment = pad_and_segment_with_inverse(x, total_segment_len)
 
                 longterm_mems, x = x[:, :num_longterm_mem_tokens], x[:, num_longterm_mem_tokens:]
@@ -277,6 +292,7 @@ class MemoryAsContextTransformer(Module):
                 x = inverse_segment(x)
 
             x = attn(x)
+
             x = ff(x)
 
         x = self.reduce_streams(x)
@@ -293,4 +309,9 @@ class MemoryAsContextTransformer(Module):
 
         x = self.norm(x)
 
-        return self.to_logits(x)
+        logits = self.to_logits(x)
+
+        if not return_loss:
+            return logits
+
+        return F.cross_entropy(rearrange(logits, 'b n l -> b l n'), labels)

{titans_pytorch-0.0.32.dist-info → titans_pytorch-0.0.34.dist-info}/METADATA

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

{titans_pytorch-0.0.32.dist-info → titans_pytorch-0.0.34.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=ohZWVhMBtpm0Iz3w5g7pD3WXSXrvhwzZvfRplwhe1Qo,8149
+titans_pytorch/mac_transformer.py,sha256=FGShQHD-dQQdQKKzvNS_jTC_FcikdqO_s3ZKOKfr_9E,8502
 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.32.dist-info/METADATA,sha256=9X9nWfgIVS-9XIeLHQY53HXSMA6rMemPfyVC2bRrJOQ,3938
-titans_pytorch-0.0.32.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.32.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.32.dist-info/RECORD,,
+titans_pytorch-0.0.34.dist-info/METADATA,sha256=CNqv_jMqk7yj15IpDn2O3jBdVe4wtrSVkht7mk0wW_E,3938
+titans_pytorch-0.0.34.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.34.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.34.dist-info/RECORD,,