x-transformers 1.32.9__py3-none-any.whl → 1.32.10__py3-none-any.whl

x_transformers/x_transformers.py

@@ -1920,7 +1920,8 @@ class TransformerWrapper(Module):
         l2norm_embed = False,
         emb_frac_gradient = 1., # GLM-130B and Cogview successfully used this, set at 0.1
         attn_z_loss_weight = 1e-4,
-        average_pool_embed = False
+        average_pool_embed = False,
+        use_cls_token = False,
     ):
         super().__init__()
 
@@ -1966,6 +1967,16 @@ class TransformerWrapper(Module):
 
         assert num_output_heads > 0
 
+        assert at_most_one_of(average_pool_embed, use_cls_token)
+
+        # classic cls token from the bert days
+
+        self.cls_token = None
+
+        if use_cls_token:
+            self.cls_token = nn.Parameter(torch.zeros(dim))
+            nn.init.normal_(self.cls_token, std = 0.02)
+
         # whether to average pool the embed (`global average pool`)
 
         self.average_pool_embed = average_pool_embed
@@ -2092,7 +2103,19 @@ class TransformerWrapper(Module):
 
         x = self.project_emb(x)
 
+        # maybe cls token
+
+        if exists(self.cls_token):
+            cls_tokens = repeat(self.cls_token, 'd -> b d', b = b)
+            x, cls_packed_shape = pack([cls_tokens, x], 'b * d')
+
+            if exists(mask):
+                mask = F.pad(mask, (1, 0), value = True)
+
+        # maybe memory / register tokens
+
         if has_memory_tokens:
+            mem_seq = x.shape[-2]
             mem_every = self.memory_tokens_interspersed_every
 
             if exists(mem_every):
@@ -2132,13 +2155,16 @@ class TransformerWrapper(Module):
             if exists(mem_every):
                 x = rearrange(x, '(b n) m d -> b (n m) d', b = b)
 
-            x = x[:, :n]
+            x = x[:, :mem_seq]
 
         # global average pool
 
         if self.average_pool_embed:
             x = masked_mean(x, mask = orig_mask, dim = 1)
 
+        if exists(self.cls_token):
+            x, _ = unpack(x, cls_packed_shape, 'b * d')
+
         # projecting to logits
 
         if not return_embeddings:

{x_transformers-1.32.9.dist-info → x_transformers-1.32.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.32.9
+Version: 1.32.10
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang

{x_transformers-1.32.9.dist-info → x_transformers-1.32.10.dist-info}/RECORD

@@ -5,11 +5,11 @@ x_transformers/continuous.py,sha256=cIVEdhfei258__ziV7kQBrJMxCel54bExBTDrO9rfCI,
 x_transformers/dpo.py,sha256=LjvWgCkqTl-UuehrzQ8nkX5guLr4whYwsmm7SKSwdls,3450
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/nonautoregressive_wrapper.py,sha256=ys_p8obc7lTeeodCqvkRKxOXQ1C9T3j5Jwr-JbVgnXk,10432
-x_transformers/x_transformers.py,sha256=8558TPHcDxWUvJYz01EdeyZl0lkHB14bzlsEMwSMPyw,77300
+x_transformers/x_transformers.py,sha256=Ao3yHjEdl-qovGo9WW8q277wBHMgFxYRfcYRf1W_hKg,78076
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-x_transformers-1.32.9.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.32.9.dist-info/METADATA,sha256=-GidCdPhcKpZ49ElbeuJUPko5LZZP_vyEodaN_P3g48,661
-x_transformers-1.32.9.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-x_transformers-1.32.9.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.32.9.dist-info/RECORD,,
+x_transformers-1.32.10.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.32.10.dist-info/METADATA,sha256=DMtabf-G60PL6axX1zSsTcWcHtzvHtKQNTxHuzOFJ4A,662
+x_transformers-1.32.10.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+x_transformers-1.32.10.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.32.10.dist-info/RECORD,,