allow for attention bias to be passed in externally

lucidrains · Oct 18, 2024 · 34a4818 · 34a4818
1 parent 7c56d23
commit 34a4818
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'x-transformers',
   packages = find_packages(exclude=['examples']),
-  version = '1.40.1',
+  version = '1.40.2',
   license='MIT',
   description = 'X-Transformers - Pytorch',
   author = 'Phil Wang',

diff --git a/x_transformers/x_transformers.py b/x_transformers/x_transformers.py
@@ -1107,6 +1107,7 @@ def forward(
         context_mask = None,
         attn_mask = None,
         rel_pos = None,
+        attn_bias = None,
         rotary_pos_emb = None,
         prev_attn = None,
         mem = None,
@@ -1237,8 +1238,8 @@ def forward(
 
         # prepare relative positional bias, if needed
 
-        attn_bias = None
         if exists(rel_pos):
+            assert not exists(attn_bias)
             attn_bias = rel_pos(i, j)
             attn_bias = pad_at_dim(attn_bias, (num_mem_kv, 0), value = 0.) # handle memory key / values