x-transformers 1.32.11__py3-none-any.whl → 1.32.14__py3-none-any.whl

x_transformers/x_transformers.py

@@ -1925,7 +1925,8 @@ class TransformerWrapper(Module):
         attn_z_loss_weight = 1e-4,
         average_pool_embed = False,
         use_cls_token = False,
-        squeeze_out_last_dim = False
+        squeeze_out_last_dim = False,
+        token_emb: TokenEmbedding | None = None,
     ):
         super().__init__()
 
@@ -1939,7 +1940,11 @@ class TransformerWrapper(Module):
         self.shift_mem_down = shift_mem_down
 
         self.l2norm_embed = l2norm_embed
-        self.token_emb = TokenEmbedding(emb_dim, num_tokens, l2norm_embed = l2norm_embed)
+
+        if not exists(token_emb):
+            token_emb = TokenEmbedding(emb_dim, num_tokens, l2norm_embed = l2norm_embed)
+
+        self.token_emb = token_emb
 
         no_abs_pos_emb = max_seq_len == 0 or not (use_abs_pos_emb and not attn_layers.disable_abs_pos_emb)
 
@@ -2184,7 +2189,7 @@ class TransformerWrapper(Module):
         # maybe squeeze out last dimension of logits
 
         if self.squeeze_out_last_dim:
-            logits = tuple(rearrange(t, '... 1 -> ...') for t in cast_tuple(logits))
+            logits = tuple((rearrange(t, '... 1 -> ...') if t.shape[-1] == 1 else t) for t in cast_tuple(logits))
 
             if not self.has_multiple_heads:
                 logits = first(logits)

{x_transformers-1.32.11.dist-info → x_transformers-1.32.14.dist-info}/METADATA

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

{x_transformers-1.32.11.dist-info → x_transformers-1.32.14.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=Z04p-xySEkTgHSaY_060M0RlF6LnkK8ko5yTLunIYf8,78520
+x_transformers/x_transformers.py,sha256=pyRQ6lb1Sx1CbjOH882tAv9UhAzsLwIeXDPBOsiRipg,78669
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-x_transformers-1.32.11.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.32.11.dist-info/METADATA,sha256=dVUc_T7ALnVvaSKHSiPKOR4Y4zAo19l0QedFQPOEbN8,662
-x_transformers-1.32.11.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-x_transformers-1.32.11.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.32.11.dist-info/RECORD,,
+x_transformers-1.32.14.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.32.14.dist-info/METADATA,sha256=ChbSwpAxqxzvSjyOlJb9yO1GnoNZk2c0ioq5F-NuHI0,662
+x_transformers-1.32.14.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+x_transformers-1.32.14.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.32.14.dist-info/RECORD,,