x-transformers 1.41.3__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):

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

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

{x_transformers-1.41.3.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=aaxMw4iJkLAh8D4W0g8EwWDfIJVPYdhpmj5T9ZFa0js,91642
+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.3.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.41.3.dist-info/METADATA,sha256=GyMdKP9ErEnGHbBqY3MwpSGSE41lK_KQbeE5n_O2Py4,689
-x_transformers-1.41.3.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-x_transformers-1.41.3.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.41.3.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,,