x-transformers 1.40.2__tar.gz → 1.40.3__tar.gz

{x_transformers-1.40.2/x_transformers.egg-info → x_transformers-1.40.3}/PKG-INFO

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

{x_transformers-1.40.2 → x_transformers-1.40.3}/README.md

@@ -525,8 +525,8 @@ model = TransformerWrapper(
         dim = 512,
         depth = 6,
         heads = 8,
-        attn_sparse_topk = 8,                   # keep only the top 8 values before attention (softmax)
-        sparse_topk_straight_through = True     # straight through the original gradients
+        attn_sparse_topk = 8,                       # keep only the top 8 values before attention (softmax)
+        attn_sparse_topk_straight_through = True    # straight through the original gradients
     )
 )
 ```

{x_transformers-1.40.2 → x_transformers-1.40.3}/setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'x-transformers',
   packages = find_packages(exclude=['examples']),
-  version = '1.40.2',
+  version = '1.40.3',
   license='MIT',
   description = 'X-Transformers - Pytorch',
   author = 'Phil Wang',

{x_transformers-1.40.2 → x_transformers-1.40.3}/x_transformers/x_transformers.py

@@ -442,10 +442,10 @@ class DynamicPositionBias(Module):
         return bias
 
 class AlibiPositionalBias(Module):
-    def __init__(self, heads, total_heads, **kwargs):
+    def __init__(self, heads, total_heads = None, **kwargs):
         super().__init__()
         self.heads = heads
-        self.total_heads = total_heads
+        self.total_heads = default(total_heads, heads)
 
         slopes = Tensor(self._get_slopes(heads))
         slopes = rearrange(slopes, 'h -> h 1 1')
@@ -1665,6 +1665,7 @@ class AttentionLayers(Module):
         cache_age = 1,
         return_hiddens = False,
         rotary_pos_emb = None,
+        attn_bias = None,
         condition = None,
         layers_execute_order: tuple[int, ...] | None = None
     ):
@@ -1818,7 +1819,7 @@ class AttentionLayers(Module):
             block = partial(block, **block_forward_kwargs)
 
             if layer_type == 'a':
-                out, inter = block(x, mask = mask, context_mask = self_attn_kv_mask, attn_mask = attn_mask, rel_pos = self.rel_pos, rotary_pos_emb = rotary_pos_emb, prev_attn = prev_attn, cache = next(iter_attn_cache, None), mem = layer_mem, mem_mask = layer_mem_mask, return_intermediates = True)
+                out, inter = block(x, mask = mask, context_mask = self_attn_kv_mask, attn_mask = attn_mask, rel_pos = self.rel_pos, rotary_pos_emb = rotary_pos_emb, prev_attn = prev_attn, cache = next(iter_attn_cache, None), mem = layer_mem, mem_mask = layer_mem_mask, attn_bias = attn_bias, return_intermediates = True)
             elif layer_type == 'c':
                 out, inter = block(x, context = context, mask = mask, context_mask = context_mask, prev_attn = prev_cross_attn, cache = next(iter_attn_cache, None), return_intermediates = True)
             elif layer_type == 'f':

{x_transformers-1.40.2 → x_transformers-1.40.3/x_transformers.egg-info}/PKG-INFO

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