x-transformers 1.42.8__py3-none-any.whl → 1.42.10__py3-none-any.whl

x_transformers/autoregressive_wrapper.py

@@ -48,7 +48,7 @@ def align_right(t, lens, pad_id = 0):
     batch_arange = torch.arange(batch, device = device, dtype = torch.long)[..., None]
     prompt_len_arange = torch.arange(seq_len, device = device, dtype = torch.long)
 
-    t = F.pad(t, (max_pad_len, 0), value = 0)
+    t = F.pad(t, (max_pad_len, 0), value = pad_id)
     offset = max_pad_len - pad_lens
 
     aligned = t[batch_arange, prompt_len_arange + offset[..., None]]

x_transformers/x_transformers.py

@@ -452,13 +452,20 @@ class DynamicPositionBias(Module):
         return bias
 
 class AlibiPositionalBias(Module):
-    def __init__(self, heads, total_heads = None, **kwargs):
+    def __init__(
+        self,
+        heads,
+        total_heads = None,
+        slopes: list[int] | None = None,
+        **kwargs
+    ):
         super().__init__()
         self.heads = heads
         self.total_heads = default(total_heads, heads)
 
-        slopes = Tensor(self._get_slopes(heads))
+        slopes = Tensor(default(slopes, self._get_slopes(heads)))
         slopes = rearrange(slopes, 'h -> h 1 1')
+
         self.register_buffer('slopes', slopes, persistent = False)
         self.register_buffer('bias', None, persistent = False)
     
@@ -487,7 +494,10 @@ class AlibiPositionalBias(Module):
         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 = -einx.subtract('... j, ... i -> ... i j', pos_j, pos_i).abs()
+
+        if bias.ndim == 3:
+            bias = rearrange(bias, 'b i j -> b 1 i j')
 
         bias = bias * self.slopes
         num_heads_unalibied = h - bias.shape[-3]
@@ -1531,8 +1541,9 @@ class AttentionLayers(Module):
         use_layerscale = False,
         layerscale_init_value = 0.,
         unet_skips = False,
-        reinject_input = False, # seen first in DEQ paper https://arxiv.org/abs/1909.01377, but later used in a number of papers trying to achieve depthwise generalization https://arxiv.org/abs/2410.03020v1
-        add_value_residual = False, # resformer from Zhou et al - https://arxiv.org/abs/2410.17897v1
+        reinject_input = False,         # seen first in DEQ paper https://arxiv.org/abs/1909.01377, but later used in a number of papers trying to achieve depthwise generalization https://arxiv.org/abs/2410.03020v1
+        add_value_residual = False,     # resformer from Zhou et al - https://arxiv.org/abs/2410.17897v1
+        rel_pos_kwargs: dict = dict(),
         **kwargs
     ):
         super().__init__()
@@ -1573,14 +1584,14 @@ class AttentionLayers(Module):
 
         if rel_pos_bias:
             assert not flash_attn, 'flash attention not compatible with t5 relative positional bias'
-            self.rel_pos = RelativePositionBias(scale = dim_head ** 0.5, causal = causal, heads = heads, num_buckets = rel_pos_num_buckets, max_distance = rel_pos_max_distance)
+            self.rel_pos = RelativePositionBias(scale = dim_head ** 0.5, causal = causal, heads = heads, num_buckets = rel_pos_num_buckets, max_distance = rel_pos_max_distance, **rel_pos_kwargs)
         elif dynamic_pos_bias:
             assert not flash_attn, 'flash attention not compatible with dynamic positional bias'
-            self.rel_pos = DynamicPositionBias(dim = dim // 4, heads = heads, log_distance = dynamic_pos_bias_log_distance, depth = dynamic_pos_bias_mlp_depth, norm = dynamic_pos_bias_norm)
+            self.rel_pos = DynamicPositionBias(dim = dim // 4, heads = heads, log_distance = dynamic_pos_bias_log_distance, depth = dynamic_pos_bias_mlp_depth, norm = dynamic_pos_bias_norm, **rel_pos_kwargs)
         elif alibi_pos_bias:
             alibi_num_heads = default(alibi_num_heads, heads)
             assert alibi_num_heads <= heads, 'number of ALiBi heads must be less than the total number of heads'
-            self.rel_pos = AlibiPositionalBias(heads = alibi_num_heads, total_heads = heads)
+            self.rel_pos = AlibiPositionalBias(heads = alibi_num_heads, total_heads = heads, **rel_pos_kwargs)
 
         assert at_most_one_of(sandwich_norm, resi_dual), 'either sandwich norm or resiDual is selected, but not both'
         assert not (not pre_norm and sandwich_norm), 'sandwich norm cannot be used when not using prenorm'

{x_transformers-1.42.8.dist-info → x_transformers-1.42.10.dist-info}/METADATA

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

{x_transformers-1.42.8.dist-info → x_transformers-1.42.10.dist-info}/RECORD

@@ -1,16 +1,16 @@
 x_transformers/__init__.py,sha256=l0dom8ZYkRzFvnDdgzDboXqrI1tKav3beVE7TN2nHko,844
 x_transformers/attend.py,sha256=SdWlV8Vp5DtpsOzAd0LRhm4VGrJf0lJCGiV2_j_CtoA,17284
-x_transformers/autoregressive_wrapper.py,sha256=DOJJCMMDOqDYKWy_IaG5IyKsXD3AW6amzfUgdAADOLY,10500
+x_transformers/autoregressive_wrapper.py,sha256=reLCno9Z9pchVU79tBF8OMo21LwSZ67KAeB83jqkyAc,10505
 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/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
-x_transformers/x_transformers.py,sha256=275B_yDHePxUvlLcMNgnCUmZ1qZEkwBrpk6IA8n-pnY,93550
+x_transformers/x_transformers.py,sha256=VxdA44EYQhVH1Rp7wreJ83I2e0Ea7VN_bFRE-iDXOI8,93833
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-1.42.8.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.42.8.dist-info/METADATA,sha256=1d2BVA6iHKpT4UzbYxw16ijAFGJT-u29zTnYtV6Lp3w,689
-x_transformers-1.42.8.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-x_transformers-1.42.8.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.42.8.dist-info/RECORD,,
+x_transformers-1.42.10.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.42.10.dist-info/METADATA,sha256=UlOzLgz1fhOIaz3eKDhcHgVI81EZySt4Ko3BiKJ7Jok,690
+x_transformers-1.42.10.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+x_transformers-1.42.10.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.42.10.dist-info/RECORD,,