x-transformers 1.27.22__py3-none-any.whl → 1.28.1__py3-none-any.whl

x_transformers/x_transformers.py

@@ -889,7 +889,15 @@ class Attention(nn.Module):
                 else:
                     input_mask = torch.cat((mem_mask, input_mask), dim = -1)
 
-        if self.num_mem_kv > 0:
+        # i, j determined for relative positional bias, excluding memory key / values
+
+        i, j = map(lambda t: t.shape[-2], (q, k))
+
+        # maybe append memory key / values
+
+        has_mem_kv = self.num_mem_kv > 0
+
+        if has_mem_kv:
             mem_k, mem_v = map(lambda t: repeat(t, 'h n d -> b h n d', b = b), (self.mem_k, self.mem_v))
 
             if self.qk_norm:
@@ -902,8 +910,6 @@ class Attention(nn.Module):
             if exists(input_mask):
                 input_mask = pad_at_dim(input_mask, (self.num_mem_kv, 0), dim = -1, value = True)
 
-        i, j = map(lambda t: t.shape[-2], (q, k))
-
         # determine masking
 
         mask_value = max_neg_value(q)
@@ -938,6 +944,11 @@ class Attention(nn.Module):
         if exists(rel_pos):
             attn_bias = rel_pos(i, j)
 
+        # append with no bias for memory key / values
+
+        if exists(attn_bias) and has_mem_kv:
+            attn_bias = pad_at_dim(attn_bias, (self.num_mem_kv, 0), value = 0.)
+
         # attention is all we need
 
         out, intermediates = self.attend(

{x_transformers-1.27.22.dist-info → x_transformers-1.28.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.27.22
+Version: 1.28.1
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang

{x_transformers-1.27.22.dist-info → x_transformers-1.28.1.dist-info}/RECORD

@@ -4,11 +4,11 @@ x_transformers/autoregressive_wrapper.py,sha256=gYKIN5Rm8dMYSTX5yHpg9sPYyZf9rsRT
 x_transformers/continuous.py,sha256=dpHK4NSMDQAJQ_N3Uj9rip0fYGXyu0QCCO_OfEdbRGs,6192
 x_transformers/dpo.py,sha256=LjvWgCkqTl-UuehrzQ8nkX5guLr4whYwsmm7SKSwdls,3450
 x_transformers/nonautoregressive_wrapper.py,sha256=ys_p8obc7lTeeodCqvkRKxOXQ1C9T3j5Jwr-JbVgnXk,10432
-x_transformers/x_transformers.py,sha256=kQhRUMGDsinzkdYcOfE1GriJ057j7D4xSjbH79FFRSE,63574
+x_transformers/x_transformers.py,sha256=F1ukKlTGhNaYmMLnm7nJy0XaUyAYwtrGGq8Gw94RPyQ,63919
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=EN3hxxleTRGYeAz6i4x3U_PrOm9TjxMF3eDhMKGx59E,8575
-x_transformers-1.27.22.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.27.22.dist-info/METADATA,sha256=RTbXIIpFRnve8FVp8vLQ4LE-9x59IV6ADnu34gGAZXA,662
-x_transformers-1.27.22.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-x_transformers-1.27.22.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.27.22.dist-info/RECORD,,
+x_transformers-1.28.1.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.28.1.dist-info/METADATA,sha256=BjW_kBCR9LCWa1CbdAkd7vMBEVYkXTgLBbNw74Tz2R4,661
+x_transformers-1.28.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+x_transformers-1.28.1.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.28.1.dist-info/RECORD,,