x-transformers 1.25.0__tar.gz → 1.25.1__tar.gz

{x-transformers-1.25.0/x_transformers.egg-info → x-transformers-1.25.1}/PKG-INFO

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

{x-transformers-1.25.0 → x-transformers-1.25.1}/setup.py

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

{x-transformers-1.25.0 → x-transformers-1.25.1}/x_transformers/x_transformers.py

@@ -1341,7 +1341,7 @@ class PrefixDecoder(AttentionLayers):
         x,
         *args,
         attn_mask = None,
-        prefix_len = None,
+        prefix_attn_len = None,
         **kwargs
     ):
         b, n, device = *x.shape[:2], x.device
@@ -1349,8 +1349,11 @@ class PrefixDecoder(AttentionLayers):
 
         forwarded_mask = ~causal_mask
 
-        if exists(prefix_len):
-            prefix_mask = torch.arange(n, device = device) < rearrange(prefix_len, 'b -> b 1 1 1')
+        if exists(prefix_attn_len):
+            if isinstance(prefix_attn_len, int):
+                prefix_attn_len = torch.full((b,), prefix_attn_len, device = device)
+
+            prefix_mask = torch.arange(n, device = device) < rearrange(prefix_attn_len, 'b -> b 1 1 1')
             forwarded_mask = forwarded_mask | prefix_mask
 
         if exists(attn_mask):

{x-transformers-1.25.0 → x-transformers-1.25.1/x_transformers.egg-info}/PKG-INFO

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