x-transformers 2.5.1__py3-none-any.whl → 2.5.3__py3-none-any.whl

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):

{x_transformers-2.5.1.dist-info → x_transformers-2.5.3.dist-info}/METADATA

@@ -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.dist-info → x_transformers-2.5.3.dist-info}/RECORD

@@ -9,10 +9,10 @@ x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
 x_transformers/up_wrapper.py,sha256=YC2LN14_7Xx9Wtiek2rtEJ_qHqdfSmKlh3d7Cgxwd80,7073
-x_transformers/x_transformers.py,sha256=l6RSwfvWxeSQWXiZVQzomyn8CyMvfmXr8oBVfb9zbBM,119679
+x_transformers/x_transformers.py,sha256=fW-AoomNCw4n2JFbZN9rZV3lKQvz_Tl6L4txUvac_9o,119993
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.5.1.dist-info/METADATA,sha256=H2y1hNleReXPTU2G1VX1vowYLQFFmaotou1QmXqJM-Q,90223
-x_transformers-2.5.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.5.1.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.5.1.dist-info/RECORD,,
+x_transformers-2.5.3.dist-info/METADATA,sha256=iR77ECuqz3O70zaZ5Mx3NwbNNal-FerMlKPlXTbv8vE,90223
+x_transformers-2.5.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.5.3.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.5.3.dist-info/RECORD,,