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

x_transformers/x_transformers.py

@@ -2038,6 +2038,9 @@ class AttentionLayers(Module):
         self.causal = causal
         self.layers = ModuleList([])
 
+        self.attn_heads = heads
+        self.attn_dim_head = dim_head
+
         # routing related
         # 1. greater than one residual stream, proposed in Hyper-Connections paper https://arxiv.org/abs/2409.19606
         # 2. integrating more than one past layer, from LIMe paper https://arxiv.org/abs/2502.09245
@@ -2758,6 +2761,9 @@ class AttentionPool(Module):
         dim_context = None,
         add_residual = False,
         depth = 1,
+        heads = 8,
+        dim_head = 64,
+        use_transformer_blocks = None,
         squeeze_output = None,
         attn_kwargs: dict = dict()
     ):
@@ -2767,13 +2773,18 @@ 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'
-            self.pooler = CrossAttender(dim = dim, cross_attn_dim_context = dim_context, depth = depth, **attn_kwargs)
+            attn_kwargs = {f'attn_{k}': v for k, v in attn_kwargs.items()}
+
+            self.pooler = CrossAttender(dim = dim, cross_attn_dim_context = dim_context, depth = depth, heads = heads, attn_dim_head = dim_head, )
         else:
-            self.pooler = Attention(dim = dim, dim_context = dim_context, **attn_kwargs)
+            self.pooler = Attention(dim = dim, dim_context = dim_context, heads = heads, dim_head = dim_head, **attn_kwargs)
 
         self.add_residual = add_residual
 
@@ -2999,7 +3010,7 @@ class TransformerWrapper(Module):
         self.attn_pool = None
 
         if attn_pool:
-            self.attn_pool = AttentionPool(dim = default(dim_pooled_tokens, dim), dim_context = dim, num_pooled_tokens = num_pooled_tokens, depth = attn_pool_depth)
+            self.attn_pool = AttentionPool(dim = default(dim_pooled_tokens, dim), dim_context = dim, num_pooled_tokens = num_pooled_tokens, depth = attn_pool_depth, heads = self.attn_layers.attn_heads, dim_head = self.attn_layers.attn_dim_head)
 
         # whether to average pool the embed (`global average pool`)
 

{x_transformers-2.5.0.dist-info → x_transformers-2.5.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.5.0
+Version: 2.5.2
 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.0.dist-info → x_transformers-2.5.2.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=NOTTbqDk5qZEY2MPpdIwJv4BvGGhXt_nIffrgQDXTf4,119346
+x_transformers/x_transformers.py,sha256=vmMrHP3hAQ9iAJlRN1pKmXOn7pD3mfh_ndtaR7LMPzU,119860
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.5.0.dist-info/METADATA,sha256=ZcxK61msWcsm8vEPG-FirnPTnwG6HNKtvy2ZoLluJHM,90223
-x_transformers-2.5.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.5.0.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.5.0.dist-info/RECORD,,
+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,,