x-transformers 2.5.1__tar.gz → 2.5.3__tar.gz

{x_transformers-2.5.1 → x_transformers-2.5.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.5.1
+Version: 2.5.3
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers

{x_transformers-2.5.1 → x_transformers-2.5.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-transformers"
-version = "2.5.1"
+version = "2.5.3"
 description = "X-Transformers"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{x_transformers-2.5.1 → x_transformers-2.5.3}/x_transformers/x_transformers.py

@@ -2763,6 +2763,7 @@ class AttentionPool(Module):
         depth = 1,
         heads = 8,
         dim_head = 64,
+        use_transformer_blocks = None,
         squeeze_output = None,
         attn_kwargs: dict = dict()
     ):
@@ -2772,9 +2773,12 @@ class AttentionPool(Module):
         squeeze_output = default(squeeze_output, False)
         assert not (squeeze_output and num_pooled_tokens > 1)
 
+        use_transformer_blocks = default(use_transformer_blocks, depth > 1)
+        assert use_transformer_blocks or depth == 1
+
         self.queries = nn.Parameter(torch.randn(num_pooled_tokens, dim) * 1e-2)
 
-        if depth > 1:
+        if use_transformer_blocks:
             assert not add_residual, 'residual already in effect when doing a full cross attention based transformer for pooling'
             attn_kwargs = {f'attn_{k}': v for k, v in attn_kwargs.items()}
 
@@ -2783,6 +2787,7 @@ class AttentionPool(Module):
             self.pooler = Attention(dim = dim, dim_context = dim_context, heads = heads, dim_head = dim_head, **attn_kwargs)
 
         self.add_residual = add_residual
+        self.squeeze_output = squeeze_output
 
     def forward(self, context, mask = None):
         batch = context.shape[0]
@@ -2794,6 +2799,9 @@ class AttentionPool(Module):
         if self.add_residual:
             pooled = pooled + queries
 
+        if self.squeeze_output:
+            pooled = rearrange(pooled, 'b 1 d -> b d')
+
         return pooled
 
 class ViTransformerWrapper(Module):