x-transformers 1.25.15__py3-none-any.whl → 1.26.1__py3-none-any.whl

x_transformers/x_transformers.py

@@ -1,5 +1,6 @@
 import math
 from random import random
+from typing import Dict
 
 import torch
 from torch import nn, einsum, Tensor
@@ -197,7 +198,7 @@ class TokenEmbedding(nn.Module):
         self.emb = nn.Embedding(num_tokens, dim)
 
     def forward(self, x):
-        token_emb = self.emb(x)
+        token_emb = self.emb(x.long())
         return l2norm(token_emb) if self.l2norm_embed else token_emb
 
 # positional embeddings
@@ -425,14 +426,15 @@ class RotaryEmbedding(nn.Module):
         self.scale_base = scale_base
         self.register_buffer('scale', scale)
 
-    @autocast(enabled = False)
-    def forward(self, seq_arange_or_len: Union[int, Tensor]):
+    def forward_from_seq_len(self, seq_len):
         device = self.inv_freq.device
 
-        if isinstance(seq_arange_or_len, int):
-            t = torch.arange(seq_arange_or_len, device = device)
-        else:
-            t = seq_arange_or_len
+        t = torch.arange(seq_len, device = device)
+        return self.forward(t)
+
+    @autocast(enabled = False)
+    def forward(self, t):
+        device = self.inv_freq.device
 
         t = t.type_as(self.inv_freq)
 
@@ -798,7 +800,8 @@ class Attention(nn.Module):
         return_intermediates = False,
         cache: Optional[Intermediates] = None,
     ):
-        b, n, _, h, kv_h, head_scale, device, has_context = *x.shape, self.heads, self.kv_heads, self.head_scale, x.device, exists(context)
+        b, n, h, kv_h, head_scale, device, has_context = x.shape[0], x.shape[1], self.heads, self.kv_heads, self.head_scale, x.device, exists(context)
+
         kv_input = default(context, x)
 
         q_input = x
@@ -1235,7 +1238,7 @@ class AttentionLayers(nn.Module):
 
         if not exists(rotary_pos_emb) and exists(self.rotary_pos_emb):
             max_rotary_emb_length = max(list(map(lambda m: (m.shape[1] if exists(m) else 0) + x.shape[1], mems)))
-            rotary_pos_emb = self.rotary_pos_emb(max_rotary_emb_length)
+            rotary_pos_emb = self.rotary_pos_emb.forward_from_seq_len(max_rotary_emb_length)
 
         # assume cached key / values
 
@@ -1361,7 +1364,7 @@ class PrefixDecoder(AttentionLayers):
         prefix_attn_len = None,
         **kwargs
     ):
-        b, n, device = *x.shape[:2], x.device
+        b, n, device = x.shape[0], x.shape[1], x.device
         causal_mask = torch.ones((n, n), device = device, dtype = torch.bool).triu(1)
 
         forwarded_mask = ~causal_mask
@@ -1462,6 +1465,7 @@ class TransformerWrapper(nn.Module):
         num_tokens,
         max_seq_len,
         attn_layers: AttentionLayers,
+        embed_num_tokens: Dict[str, int] = dict(),
         emb_dim = None,
         max_mem_len = 0,
         shift_mem_down = 0,
@@ -1491,14 +1495,23 @@ class TransformerWrapper(nn.Module):
         self.l2norm_embed = l2norm_embed
         self.token_emb = TokenEmbedding(emb_dim, num_tokens, l2norm_embed = l2norm_embed)
 
-        if not (use_abs_pos_emb and not attn_layers.has_pos_emb):
+        if max_seq_len == 0 or not (use_abs_pos_emb and not attn_layers.has_pos_emb):
             self.pos_emb = always(0)
         elif scaled_sinu_pos_emb:
             self.pos_emb = ScaledSinusoidalEmbedding(emb_dim)
         else:
             self.pos_emb = AbsolutePositionalEmbedding(emb_dim, max_seq_len, l2norm_embed = l2norm_embed)
 
-        self.emb_frac_gradient = emb_frac_gradient # fraction of the gradient that should go to the embedding, https://arxiv.org/abs/2105.13290
+        # additional embeddings - say type embedding from BERT
+
+        self.embeds = None
+
+        if len(embed_num_tokens) > 0:
+            self.embeds = nn.ModuleDict({f'{name}_embed': nn.Embedding(num_tokens, emb_dim) for name, num_tokens in embed_num_tokens.items()})
+
+        # fraction of the gradient that should go to the embedding, https://arxiv.org/abs/2105.13290
+
+        self.emb_frac_gradient = emb_frac_gradient
 
         self.post_emb_norm = nn.LayerNorm(emb_dim) if post_emb_norm else nn.Identity()
         self.emb_dropout = nn.Dropout(emb_dropout)
@@ -1546,6 +1559,7 @@ class TransformerWrapper(nn.Module):
         pos = None,
         prepend_embeds = None,
         prepend_mask = None,
+        embed_ids: Dict[str, Tensor] = None,
         sum_embeds = None,
         return_attn_z_loss = False,
         attn_z_loss_weight = 1e-4,
@@ -1553,7 +1567,7 @@ class TransformerWrapper(nn.Module):
         cache: Optional[LayerIntermediates] = None,
         **kwargs
     ):
-        b, n, device, num_mems, has_memory_tokens, emb_frac_gradient = *x.shape, x.device, self.num_memory_tokens, self.num_memory_tokens > 0, self.emb_frac_gradient
+        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
 
         # absolute positional embedding
@@ -1562,6 +1576,17 @@ class TransformerWrapper(nn.Module):
         pos_emb = self.pos_emb(x, pos = pos, seq_start_pos = seq_start_pos) if not external_pos_emb else pos
         x = self.token_emb(x) + pos_emb
 
+        # add additional embeddings
+
+        if exists(self.embeds) and exists(embed_ids):
+            for name, embed_id in embed_ids.items():
+                embed_key = f'{name}_embed'
+
+                assert embed_key in self.embeds
+                embed = self.embeds[embed_key](embed_id)
+
+                x = x + embed
+
         # for summing embeddings passed externally - needs this for self-conditioning in non-autoregressive training
 
         if exists(sum_embeds):

{x_transformers-1.25.15.dist-info → x_transformers-1.26.1.dist-info}/METADATA

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

{x_transformers-1.25.15.dist-info → x_transformers-1.26.1.dist-info}/RECORD

@@ -3,11 +3,11 @@ x_transformers/attend.py,sha256=MFl_FbgPsm9mziZPTi_s8QbxASETwbGeciMH8sUIwT8,1018
 x_transformers/autoregressive_wrapper.py,sha256=2mzOq_rl_vevgrxCDncBlVJOAJGS-XGm-iBKJqMjj_c,9041
 x_transformers/continuous.py,sha256=ixfgi2_zpGN03SX_STXFkNYEOAkgwVIxuS53QgDCx-g,6026
 x_transformers/nonautoregressive_wrapper.py,sha256=AQLE4rA_Kh8VNoe9OzpwyeWson34sRkhks4dn4seNjI,10414
-x_transformers/x_transformers.py,sha256=5T5fOdveqe4vwqGRrNZqJ053pmtw-dJ12AV0nNaWLRc,60814
+x_transformers/x_transformers.py,sha256=k-_pU9lICUYnCumwimPv3VaaxjpKOeFKDdKEWgnvslc,61597
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=lS9W_E_RskPQAqVZkPiUzbByoW1Ajsw_phsikA3JDAg,8139
-x_transformers-1.25.15.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.25.15.dist-info/METADATA,sha256=npEGmCXBOdJv7S8odHPNOrjTOvFgSxNFjx1377QtYn8,662
-x_transformers-1.25.15.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-x_transformers-1.25.15.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.25.15.dist-info/RECORD,,
+x_transformers-1.26.1.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.26.1.dist-info/METADATA,sha256=mJsZmH-mfNGzd7eGYglFH6GlNzGUTeugQHLcH8oDecU,661
+x_transformers-1.26.1.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+x_transformers-1.26.1.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.26.1.dist-info/RECORD,,