x-transformers 1.32.2__py3-none-any.whl → 1.32.4__py3-none-any.whl

x_transformers/x_transformers.py

@@ -1899,6 +1899,7 @@ class TransformerWrapper(Module):
         memory_tokens_interspersed_every = None,
         tie_embedding = False,
         logits_dim = None,
+        return_only_embed = False,
         num_output_heads = 1,
         use_abs_pos_emb = True,
         scaled_sinu_pos_emb = False,
@@ -1948,13 +1949,17 @@ class TransformerWrapper(Module):
 
         self.init_()
 
+        assert num_output_heads > 0
+
         # output head, usually to logits of num_tokens
 
         logits_dim = default(logits_dim, num_tokens)
 
         self.has_multiple_heads = False
 
-        if tie_embedding:
+        if return_only_embed:
+            self.to_logits = None
+        elif tie_embedding:
             self.to_logits = lambda t: t @ self.token_emb.emb.weight.t()
         elif num_output_heads > 1:
             self.has_multiple_heads = True
@@ -2008,7 +2013,9 @@ class TransformerWrapper(Module):
         **kwargs
     ):
         b, n, device, num_mems, has_memory_tokens, emb_frac_gradient = x.shape[0], x.shape[1], x.device, self.num_memory_tokens, self.num_memory_tokens > 0, self.emb_frac_gradient
+
         return_hiddens = return_mems | return_attn | return_intermediates | return_attn_z_loss
+        return_embeddings = return_embeddings | (not exists(self.to_logits))
 
         # absolute positional embedding
 
@@ -2018,6 +2025,8 @@ class TransformerWrapper(Module):
 
         # add additional embeddings
 
+        assert not (exists(self.embeds) ^ (len(embed_ids) > 0)), '`embed_num_tokens` must be defined on `TransformerWrapper`'
+
         if exists(self.embeds):
             assert len(embed_ids) == len(self.embeds)
 

{x_transformers-1.32.2.dist-info → x_transformers-1.32.4.dist-info}/METADATA

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

{x_transformers-1.32.2.dist-info → x_transformers-1.32.4.dist-info}/RECORD

@@ -4,11 +4,11 @@ x_transformers/autoregressive_wrapper.py,sha256=uX8Mb0zLsQrZECt_9UGt35g7tC05Rk3n
 x_transformers/continuous.py,sha256=WO52n9lFAXv5-SGadi2cApGF8dkouN8QSTEOuC7erj8,6180
 x_transformers/dpo.py,sha256=LjvWgCkqTl-UuehrzQ8nkX5guLr4whYwsmm7SKSwdls,3450
 x_transformers/nonautoregressive_wrapper.py,sha256=ys_p8obc7lTeeodCqvkRKxOXQ1C9T3j5Jwr-JbVgnXk,10432
-x_transformers/x_transformers.py,sha256=1QG7zUe89h1R5VDMoKEAkvdRRDkzQ7h6npkqblxxR6g,76312
+x_transformers/x_transformers.py,sha256=o7At5F35Paih1a_-rxqcIP1n4B-ARVJ_ZL2QkOnTnSQ,76655
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-x_transformers-1.32.2.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.32.2.dist-info/METADATA,sha256=U0Kh4e7UiL-0hLDZb0P3McdvTnzTeFyVwtoXFffzQ-M,661
-x_transformers-1.32.2.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-x_transformers-1.32.2.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.32.2.dist-info/RECORD,,
+x_transformers-1.32.4.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.32.4.dist-info/METADATA,sha256=rAX2cTBnI50T0Tsa00KgBTz6lzV36ACppE6H2WPLZI4,661
+x_transformers-1.32.4.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+x_transformers-1.32.4.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.32.4.dist-info/RECORD,,