x-transformers 2.3.26__py3-none-any.whl → 2.4.0__py3-none-any.whl

x_transformers/x_transformers.py

@@ -47,6 +47,8 @@ class LayerIntermediates:
     layer_hiddens:      list[Tensor] | None = None
     attn_z_loss:        Tensor | None = None
     mems:               Tensor | None = None
+    last_layer_hiddens: Tensor | None = None
+    attn_pooled_tokens: Tensor | None = None
     memory_tokens:      Tensor | None = None
     logit_entropies:    Tensor | None = None
     cache_length:       int = 0
@@ -1926,7 +1928,7 @@ class Attention(Module):
 
         out = maybe(self.sublayer_dropout)(out)
 
-        if exists(mask):
+        if exists(mask) and not exists(cache):
             out = einx.where('b n, b n d, -> b n d', mask, out, 0.)
 
         if not return_intermediates:
@@ -2484,7 +2486,7 @@ class AttentionLayers(Module):
         attn_cache = []
 
         if exists(cache):
-            assert self.causal and not any([*map(exists, (mask, attn_mask))])
+            assert self.causal and not exists(attn_mask)
 
             prev_cache_length = cache.cache_length
 
@@ -2848,6 +2850,9 @@ class TransformerWrapper(Module):
         average_pool_embed = False,
         use_cls_token = False,
         num_cls_tokens = 1,
+        attn_pool = False,
+        num_attn_pool_queries = 1,
+        dim_attn_pool_query = None,
         squeeze_out_last_dim = False,
         token_emb: TokenEmbedding | None = None,
         mixture_of_softmax = False,
@@ -2927,6 +2932,10 @@ class TransformerWrapper(Module):
 
         self.train_max_recycle_steps = train_max_recycle_steps
 
+        # either cls token or attn pool, but not both
+
+        assert not (use_cls_token and attn_pool)
+
         # classic cls token from the bert days
 
         self.cls_token = None
@@ -2935,6 +2944,16 @@ class TransformerWrapper(Module):
             self.cls_token = nn.Parameter(torch.zeros(num_cls_tokens, dim))
             nn.init.normal_(self.cls_token, std = 0.02)
 
+        # attn pool
+
+        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)
+
         # whether to average pool the embed (`global average pool`)
 
         self.average_pool_embed = average_pool_embed
@@ -3222,14 +3241,37 @@ class TransformerWrapper(Module):
 
             x = x[:, :mem_seq]
 
+        # store last layer hiddens, for access in case of cls token or attention pooling
+
+        intermediates.last_layer_hiddens = x
+
         # global average pool
 
         if self.average_pool_embed:
             x = masked_mean(x, mask = orig_mask, dim = 1)
 
+
+        # cls token(s)
+
         if exists(self.cls_token):
-            x, _ = unpack(x, cls_packed_shape, 'b * d')
-            x = x.squeeze(1)  # Remove sequence dimension if num_cls_tokens=1 to keep previous behavior
+            x, last_layer_hiddens = unpack(x, cls_packed_shape, 'b * d')
+
+            intermediates.last_layer_hiddens = last_layer_hiddens
+
+            if x.shape[1] == 1:
+                x = rearrange(x, 'b 1 d -> b d')  # Remove sequence dimension if num_cls_tokens=1 to keep previous behavior
+
+        # attention pool
+
+        if exists(self.attn_pool):
+            queries = repeat(self.attn_pool_queries, 'n d -> b n d', b = x.shape[0])
+
+            attn_pooled_tokens = self.attn_pool(queries, context = x, context_mask = mask)
+
+            if attn_pooled_tokens.shape[1] == 1:
+                attn_pooled_tokens = rearrange(attn_pooled_tokens, 'b 1 d -> b d')
+
+            intermediates.attn_pooled_tokens = attn_pooled_tokens
 
         # handle expansion to mixture if needed (for mixture of softmax)
 

{x_transformers-2.3.26.dist-info → x_transformers-2.4.0.dist-info}/METADATA

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

{x_transformers-2.3.26.dist-info → x_transformers-2.4.0.dist-info}/RECORD

@@ -8,10 +8,10 @@ x_transformers/entropy_based_tokenizer.py,sha256=F2lO8-v3aLIcVDVNhu7RR-UtRdlmaaY
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
-x_transformers/x_transformers.py,sha256=7phSZvP1_SDRIkVMwVR4cz1dFU2UlR2Wf1HJHEQlcQg,116222
+x_transformers/x_transformers.py,sha256=IelVhLUuDmRnv6zXlQNvwUluW2RqVQQE2vYKCqctJyY,117583
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.3.26.dist-info/METADATA,sha256=Qc0zIph59FLOC0GPGIe41M6P1SD_lljzKg5ytoMyPAI,89897
-x_transformers-2.3.26.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.3.26.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.3.26.dist-info/RECORD,,
+x_transformers-2.4.0.dist-info/METADATA,sha256=RyKkjmTnjbGUHA4EL-znJCPR17VF6i7ebvvgMKpTXVY,89896
+x_transformers-2.4.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.4.0.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.4.0.dist-info/RECORD,,