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

x_transformers/x_transformers.py

@@ -2787,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]
@@ -2798,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.2.dist-info → x_transformers-2.5.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.5.2
+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.2.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=vmMrHP3hAQ9iAJlRN1pKmXOn7pD3mfh_ndtaR7LMPzU,119860
+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.2.dist-info/METADATA,sha256=yeferX_PJIv0Lxs36vZSV7Z2w9ol4udiUAON95hP_bY,90223
-x_transformers-2.5.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.5.2.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.5.2.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,,