x-transformers 1.32.5__py3-none-any.whl → 1.32.7__py3-none-any.whl

x_transformers/multi_input.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import torch
 from torch import nn, Tensor
 from torch.nn import Module, ModuleDict
@@ -15,7 +17,7 @@ from x_transformers.x_transformers import (
     LayerNorm,
     always,
     pad_at_dim,
-    is_empty
+    is_empty,
 )
 
 # helper functions
@@ -43,10 +45,7 @@ class MultiInputTransformerWrapper(Module):
         post_emb_norm = False,
         num_memory_tokens = None,
         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,
         emb_frac_gradient = 1., # GLM-130B and Cogview successfully used this, set at 0.1
@@ -85,23 +84,12 @@ class MultiInputTransformerWrapper(Module):
         self.project_emb = nn.Linear(emb_dim, dim) if emb_dim != dim else nn.Identity()
         self.attn_layers = attn_layers
 
-        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 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
-            self.to_logits = ModuleList([nn.Linear(dim, logits_dim, bias = False) for _ in range(num_output_heads)])
         else:
-            self.to_logits = nn.Linear(dim, logits_dim, bias = False)
+            self.to_logits = ModuleDict({name: nn.Linear(dim, logits_dim, bias = False) for name, logits_dim in num_tokens.items()})
 
         # memory tokens (like [cls]) from Memory Transformers paper
 
@@ -251,10 +239,7 @@ class MultiInputTransformerWrapper(Module):
         # projecting to logits
 
         if not return_embeddings:
-            if self.has_multiple_heads:
-                logits = tuple(fn(x) for fn in self.to_logits)
-            else:
-                logits = self.to_logits(x)
+            logits = {name: fn(x) for name, fn in self.to_logits.items()}
 
         # different returns
 

{x_transformers-1.32.5.dist-info → x_transformers-1.32.7.dist-info}/METADATA

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

{x_transformers-1.32.5.dist-info → x_transformers-1.32.7.dist-info}/RECORD

@@ -3,13 +3,13 @@ x_transformers/attend.py,sha256=MI-m91wumBFqFqr_KK9MLgsLk_vPeaVbFMyDr_mWdmY,1134
 x_transformers/autoregressive_wrapper.py,sha256=uX8Mb0zLsQrZECt_9UGt35g7tC05Rk3nPqO6xp2FFCc,9619
 x_transformers/continuous.py,sha256=WO52n9lFAXv5-SGadi2cApGF8dkouN8QSTEOuC7erj8,6180
 x_transformers/dpo.py,sha256=LjvWgCkqTl-UuehrzQ8nkX5guLr4whYwsmm7SKSwdls,3450
-x_transformers/multi_input.py,sha256=S4hJffAhN7Ubbd92_9YmZ0VOzdu1u-1e2IgWZxYz9BU,9758
+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=5DHbYgx0RPg9QHvfBs2qHWrtn4Jji-q0d1MRBbcRPR8,76696
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-x_transformers-1.32.5.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.32.5.dist-info/METADATA,sha256=HSujd9BjY5lAIY5QxDqf4JDecEOXrvWQoe2zjpRaGG0,661
-x_transformers-1.32.5.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-x_transformers-1.32.5.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.32.5.dist-info/RECORD,,
+x_transformers-1.32.7.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.32.7.dist-info/METADATA,sha256=25J9CJ3OxsR_SZkvubPhyjSN-NmvU_yVVQHNMFzoKVg,661
+x_transformers-1.32.7.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+x_transformers-1.32.7.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.32.7.dist-info/RECORD,,