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

x_transformers/__init__.py

@@ -4,12 +4,13 @@ from x_transformers.x_transformers import (
     Decoder,
     PrefixDecoder,
     CrossAttender,
+    AttentionPool,
     Attention,
     FeedForward,
     RMSNorm,
     AdaptiveRMSNorm,
     TransformerWrapper,
-    ViTransformerWrapper
+    ViTransformerWrapper,
 )
 
 from x_transformers.autoregressive_wrapper import AutoregressiveWrapper

x_transformers/x_transformers.py

@@ -51,6 +51,7 @@ class LayerIntermediates:
     attn_pooled_tokens: Tensor | None = None
     memory_tokens:      Tensor | None = None
     logit_entropies:    Tensor | None = None
+    logits:             Tensor | None = None
     cache_length:       int = 0
 
 LinearNoBias = partial(nn.Linear, bias = False)
@@ -2037,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
@@ -2749,6 +2753,49 @@ class CrossAttender(AttentionLayers):
     def __init__(self, **kwargs):
         super().__init__(cross_attend = True, only_cross = True, **kwargs)
 
+class AttentionPool(Module):
+    def __init__(
+        self,
+        dim,
+        num_pooled_tokens = 1,
+        dim_context = None,
+        add_residual = False,
+        depth = 1,
+        heads = 8,
+        dim_head = 64,
+        squeeze_output = None,
+        attn_kwargs: dict = dict()
+    ):
+        super().__init__()
+        dim_context = default(dim_context, dim)
+
+        squeeze_output = default(squeeze_output, False)
+        assert not (squeeze_output and num_pooled_tokens > 1)
+
+        self.queries = nn.Parameter(torch.randn(num_pooled_tokens, dim) * 1e-2)
+
+        if depth > 1:
+            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()}
+
+            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, heads = heads, dim_head = dim_head, **attn_kwargs)
+
+        self.add_residual = add_residual
+
+    def forward(self, context, mask = None):
+        batch = context.shape[0]
+
+        queries = repeat(self.queries, 'n d -> b n d', b = batch)
+
+        pooled = self.pooler(queries, context, context_mask = mask)
+
+        if self.add_residual:
+            pooled = pooled + queries
+
+        return pooled
+
 class ViTransformerWrapper(Module):
     def __init__(
         self,
@@ -2860,8 +2907,9 @@ class TransformerWrapper(Module):
         use_cls_token = False,
         num_cls_tokens = 1,
         attn_pool = False,
-        num_attn_pool_queries = 1,
-        dim_attn_pool_query = None,
+        num_pooled_tokens = 1,
+        attn_pool_depth = 1,
+        dim_pooled_tokens = None,
         squeeze_out_last_dim = False,
         token_emb: TokenEmbedding | None = None,
         mixture_of_softmax = False,
@@ -2958,10 +3006,7 @@ class TransformerWrapper(Module):
         self.attn_pool = None
 
         if attn_pool:
-            self.attn_pool = Attention(dim = default(dim_attn_pool_query, dim), dim_context = dim)
-
-            self.attn_pool_queries = nn.Parameter(torch.zeros(num_attn_pool_queries, dim))
-            nn.init.normal_(self.attn_pool_queries, std = 0.02)
+            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`)
 
@@ -3259,7 +3304,6 @@ class TransformerWrapper(Module):
         if self.average_pool_embed:
             x = masked_mean(x, mask = orig_mask, dim = 1)
 
-
         # cls token(s)
 
         if exists(self.cls_token):
@@ -3272,13 +3316,15 @@ class TransformerWrapper(Module):
 
         # attention pool
 
-        if exists(self.attn_pool) and return_intermediates:
-            queries = repeat(self.attn_pool_queries, 'n d -> b n d', b = x.shape[0])
+        is_encoder = not self.attn_layers.causal
+        return_pooled_tokens = exists(self.attn_pool) and is_encoder
 
-            attn_pooled_tokens = self.attn_pool(queries, context = x, context_mask = mask)
+        if (
+            exists(self.attn_pool) and
+            (return_intermediates or is_encoder) # in a new paper, they use attention pooling on decoder - so we'll default to returning pooled tokens if encoder, but for decoder, they must set `return_intermediates`
+        ):
 
-            if attn_pooled_tokens.shape[1] == 1:
-                attn_pooled_tokens = rearrange(attn_pooled_tokens, 'b 1 d -> b d')
+            attn_pooled_tokens = self.attn_pool(x, mask = mask)
 
             intermediates.attn_pooled_tokens = attn_pooled_tokens
 
@@ -3327,6 +3373,9 @@ class TransformerWrapper(Module):
             out = (x, intermediates)
         elif return_embeddings:
             out = x
+        elif return_pooled_tokens:
+            intermediates.logits = logits
+            out = attn_pooled_tokens
         else:
             out = logits
 

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

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

{x_transformers-2.4.14.dist-info → x_transformers-2.5.1.dist-info}/RECORD

@@ -1,4 +1,4 @@
-x_transformers/__init__.py,sha256=h3I2ejobgEdy8H7NgV-rP8UaBCnd16-MysvDXH9GMEA,985
+x_transformers/__init__.py,sha256=aVuhUU0572TJHW88BVc4yA2tla0Zb8l3NH7W4RZ1AEs,1005
 x_transformers/attend.py,sha256=Ax34Rw56xXAWfFPqtZ_B8iKEW2EfQdbVoc9uFjfeNjA,17404
 x_transformers/autoregressive_wrapper.py,sha256=y798kS9_VvPOY_5Ilits_64aXNqYvGuilsky1y07ryE,17834
 x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTNY__GGZEKU,13374
@@ -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=bNp6hWuuqn7x5yKFfYocvu3X1YCjpfwrWMh-kAanS48,117906
+x_transformers/x_transformers.py,sha256=l6RSwfvWxeSQWXiZVQzomyn8CyMvfmXr8oBVfb9zbBM,119679
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.4.14.dist-info/METADATA,sha256=KScRZIcmRXCv8NnhzQ3Uzo9uHE2oI51chzj78Wh_OVo,90224
-x_transformers-2.4.14.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.4.14.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.4.14.dist-info/RECORD,,
+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,,