x-transformers 1.41.2__py3-none-any.whl → 1.41.4__py3-none-any.whl

x_transformers/attend.py

@@ -480,6 +480,7 @@ class Attend(Module):
             sim = sim + self.pre_softmax_talking_heads(sim)
 
         if exists(attn_bias):
+            print(attn_bias.shape)
             sim = sim + attn_bias
 
         if self.softclamp_logits:

x_transformers/x_transformers.py

@@ -455,11 +455,9 @@ class AlibiPositionalBias(Module):
         self.register_buffer('slopes', slopes, persistent = False)
         self.register_buffer('bias', None, persistent = False)
     
-    def get_bias(self, i, j, device):
-        seq_arange = torch.arange(j - i, j, device = device)
-        context_arange = torch.arange(j, device = device)
-        bias = -torch.abs(einx.subtract('j, i -> 1 i j', context_arange, seq_arange))
-        return bias
+    @property
+    def device(self):
+        return next(self.buffers()).device
 
     @staticmethod
     def _get_slopes(heads):
@@ -474,9 +472,21 @@ class AlibiPositionalBias(Module):
         closest_power_of_2 = 2 ** math.floor(math.log2(heads))
         return get_slopes_power_of_2(closest_power_of_2) + get_slopes_power_of_2(2 * closest_power_of_2)[0::2][:heads-closest_power_of_2]
 
-    @property
-    def device(self):
-        return next(self.buffers()).device
+    def forward_custom_pos(
+        self,
+        pos_i: Tensor,
+        pos_j: Tensor | None = None
+    ):
+        h, device = self.total_heads, self.device
+
+        pos_j = default(pos_j, pos_i)
+        bias = -einx.subtract('... j, ... i -> ... 1 i j', pos_j, pos_i).abs()
+
+        bias = bias * self.slopes
+        num_heads_unalibied = h - bias.shape[-3]
+        bias = pad_at_dim(bias, (0, num_heads_unalibied), dim = -3)
+
+        return bias
 
     def forward(self, i, j):
         h, device = self.total_heads, self.device
@@ -484,13 +494,15 @@ class AlibiPositionalBias(Module):
         if exists(self.bias) and self.bias.shape[-1] >= j and self.bias.shape[-2] >= i:
             return self.bias[..., -i:, -j:]
 
-        bias = self.get_bias(i, j, device)
+        seq_arange = torch.arange(j - i, j, device = device)
+        context_arange = torch.arange(j, device = device)
+        bias = -einx.subtract('j, i -> 1 i j', context_arange, seq_arange).abs()
+
         bias = bias * self.slopes
+        num_heads_unalibied = h - bias.shape[-3]
+        bias = pad_at_dim(bias, (0, num_heads_unalibied), dim = -3)
 
-        num_heads_unalibied = h - bias.shape[0]
-        bias = pad_at_dim(bias, (0, num_heads_unalibied), dim = 0)
         self.register_buffer('bias', bias, persistent = False)
-
         return self.bias
 
 class DataDependentAlibi(Module):
@@ -499,7 +511,9 @@ class DataDependentAlibi(Module):
     def __init__(
         self,
         dim,
-        heads
+        heads,
+        bias_init = 5.,
+        post_log_scale = 1.
     ):
         super().__init__()
 
@@ -511,22 +525,24 @@ class DataDependentAlibi(Module):
             nn.LogSigmoid()
         )
 
-        nn.init.constant_(linear.bias, 5.)
+        nn.init.constant_(linear.bias, bias_init)
+        self.post_log_scale = post_log_scale
 
     def forward(self, x):
-        forget_gates = self.to_forget_gates(x)
+        forget_gates = self.to_forget_gates(x) * self.post_log_scale
         forget_gates = forget_gates.cumsum(dim = -1)
         forget_gates = einx.subtract('b h i, b h j -> b h i j', forget_gates, forget_gates)
         return forget_gates
 
 class PerRowDataDependentAlibi(Module):
-    """ same as data dependent alibi from forgetting transformer, but the forgetting gates are also derived by a queris and keys with a small head dimension """
+    """ same as data dependent alibi from forgetting transformer, but the forgetting gates are also derived by a queries and keys with a small head dimension """
 
     def __init__(
         self,
         dim,
         heads,
-        dim_head = 8
+        dim_head = 8,
+        post_log_scale = 1.
     ):
         super().__init__()
         self.scale = dim_head ** -0.5
@@ -535,14 +551,16 @@ class PerRowDataDependentAlibi(Module):
 
         self.to_forget_gates = nn.Sequential(
             linear,
-            Rearrange('b n (kv h d) -> kv b h n d', kv = 2, d = dim_head)
+            Rearrange('b n (qk h d) -> qk b h n d', qk = 2, d = dim_head)
         )
 
+        self.post_log_scale = post_log_scale
+
     def forward(self, x):
         q, k = self.to_forget_gates(x)
         forget_gates = einsum('... i d, ... j d -> ... i j', q, k) * self.scale
 
-        forget_gates = F.logsigmoid(forget_gates)
+        forget_gates = F.logsigmoid(forget_gates) * self.post_log_scale
 
         # mask out upper triangle + diagonal
 
@@ -1010,6 +1028,7 @@ class Attention(Module):
         data_dependent_alibi = False,
         data_dependent_alibi_per_row = False,
         data_dependent_alibi_per_row_dim_head = 8,
+        data_dependent_alibi_kwargs: dict = dict(),
         use_cope = False,
         cope_max_pos = 16,
         cope_soft_onehot_pos = False,
@@ -1127,7 +1146,7 @@ class Attention(Module):
             if data_dependent_alibi_per_row:
                 dda_kwargs.update(dim_head = data_dependent_alibi_per_row_dim_head)
 
-            self.data_dependent_alibi = dda_klass(**dda_kwargs)
+            self.data_dependent_alibi = dda_klass(**dda_kwargs, **data_dependent_alibi_kwargs)
 
         # attend class - includes core attention algorithm + talking heads
 

{x_transformers-1.41.2.dist-info → x_transformers-1.41.4.dist-info}/METADATA

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

{x_transformers-1.41.2.dist-info → x_transformers-1.41.4.dist-info}/RECORD

@@ -1,15 +1,15 @@
 x_transformers/__init__.py,sha256=-MkQrSc37cTVDX7AOykxunYnqVtFlQ7lb0Cse5dsGWU,793
-x_transformers/attend.py,sha256=SdWlV8Vp5DtpsOzAd0LRhm4VGrJf0lJCGiV2_j_CtoA,17284
+x_transformers/attend.py,sha256=rByHtOfuCO0br69rOB7oFsHoHrsAefZErE2FKM86q7k,17319
 x_transformers/autoregressive_wrapper.py,sha256=DOJJCMMDOqDYKWy_IaG5IyKsXD3AW6amzfUgdAADOLY,10500
 x_transformers/continuous.py,sha256=cIVEdhfei258__ziV7kQBrJMxCel54bExBTDrO9rfCI,6450
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
-x_transformers/x_transformers.py,sha256=5yoN96mTpw3qygZg2plVed5lTOs9p0REtt50lOlllHU,91334
+x_transformers/x_transformers.py,sha256=UhIbFPXjdQsbFBDHVGmV81LGHTD5qwbusDc5kl3F2A4,91987
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-1.41.2.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.41.2.dist-info/METADATA,sha256=84WHjWdVHIxkHdHxMEKK5OEssLU80cdj4aGTJQ7SpO4,689
-x_transformers-1.41.2.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-x_transformers-1.41.2.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.41.2.dist-info/RECORD,,
+x_transformers-1.41.4.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.41.4.dist-info/METADATA,sha256=BWpwY2VILr3LAvjqQjDBm28jpS3RZWvmLrghAa7HeuU,689
+x_transformers-1.41.4.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+x_transformers-1.41.4.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.41.4.dist-info/RECORD,,