x-transformers 1.42.6__py3-none-any.whl → 1.42.8__py3-none-any.whl

x_transformers/x_transformers.py

@@ -238,6 +238,13 @@ class TokenEmbedding(Module):
         token_emb = self.emb(x.long())
         return l2norm(token_emb) if self.l2norm_embed else token_emb
 
+    def init_(self):
+        if self.l2norm_embed:
+            nn.init.normal_(self.emb.weight, std=1e-5)
+            return
+        nn.init.kaiming_normal_(self.emb.weight)
+
+
 # positional embeddings
 
 class AbsolutePositionalEmbedding(Module):
@@ -1246,6 +1253,7 @@ class Attention(Module):
         rel_pos = None,
         attn_bias = None,
         rotary_pos_emb = None,
+        pos = None, # for custom alibi positions
         prev_attn = None,
         mem = None,
         mem_mask = None,
@@ -1392,7 +1400,14 @@ class Attention(Module):
 
         if exists(rel_pos):
             assert not exists(attn_bias)
-            attn_bias = rel_pos(i, j)
+
+            if exists(pos):
+                assert isinstance(rel_pos, AlibiPositionalBias), 'only alibi allowed for custom positions at the moment'
+                # allow for custom positions to be passed in
+                attn_bias = rel_pos.forward_custom_pos(pos)
+            else:
+                attn_bias = rel_pos(i, j)
+
             attn_bias = pad_at_dim(attn_bias, (num_mem_kv, 0), value = 0.) # handle memory key / values
 
         # prepare data dependent alibi from forgetting transformers paper, if needed
@@ -1843,6 +1858,7 @@ class AttentionLayers(Module):
         cache_age = 1,
         return_hiddens = False,
         rotary_pos_emb = None,
+        pos = None,
         attn_bias = None,
         condition = None,
         in_attn_cond = None, # https://arxiv.org/abs/2105.04090
@@ -1906,7 +1922,9 @@ class AttentionLayers(Module):
             maybe_mem = mems[0] # todo - handle edge case where different layers get different memory lengths. don't think this will ever come up but who knows
             mem_len = maybe_mem.shape[1] if exists(maybe_mem) else 0
 
-            pos = torch.arange(x.shape[1] + mem_len, device = x.device) - mem_len
+            if not exists(pos):
+                pos = torch.arange(x.shape[1] + mem_len, device = x.device) - mem_len
+
             rotary_pos_emb = self.rotary_pos_emb(pos)
 
         # assume cached key / values
@@ -2030,7 +2048,7 @@ class AttentionLayers(Module):
             # forward depending on layer type
 
             if layer_type == 'a':
-                out, inter = block(x, mask = mask, context_mask = self_attn_kv_mask, attn_mask = attn_mask, rel_pos = self.rel_pos, rotary_pos_emb = rotary_pos_emb, prev_attn = prev_attn, cache = next(iter_attn_cache, None), mem = layer_mem, mem_mask = layer_mem_mask, attn_bias = attn_bias, value_residual = maybe_self_attn_value_residual, return_intermediates = True)
+                out, inter = block(x, mask = mask, context_mask = self_attn_kv_mask, attn_mask = attn_mask, rel_pos = self.rel_pos, pos = pos, rotary_pos_emb = rotary_pos_emb, prev_attn = prev_attn, cache = next(iter_attn_cache, None), mem = layer_mem, mem_mask = layer_mem_mask, attn_bias = attn_bias, value_residual = maybe_self_attn_value_residual, return_intermediates = True)
             elif layer_type == 'c':
                 out, inter = block(x, context = context, mask = mask, context_mask = context_mask, prev_attn = prev_cross_attn, cache = next(iter_attn_cache, None), value_residual = maybe_cross_attn_value_residual, return_intermediates = True)
             elif layer_type == 'f':
@@ -2250,7 +2268,8 @@ class TransformerWrapper(Module):
         token_emb: TokenEmbedding | None = None,
         mixture_of_softmax = False,
         mixture_of_softmax_k = 4,
-        sigsoftmax_logits = False
+        sigsoftmax_logits = False,
+        to_logits: Module | None = None,
     ):
         super().__init__()
 
@@ -2352,11 +2371,12 @@ class TransformerWrapper(Module):
         if return_only_embed:
             self.to_logits = None
         elif tie_embedding:
+            assert isinstance(token_emb, TokenEmbedding), 'can only tie embedding if using `TokenEmbedding`'
             self.to_logits = lambda t: t @ self.token_emb.emb.weight.t()
         elif num_output_heads > 1:
             self.to_logits = ModuleList([LinearNoBias(dim, logits_dim) for _ in range(num_output_heads)])
         else:
-            self.to_logits = LinearNoBias(dim, logits_dim)
+            self.to_logits = LinearNoBias(dim, logits_dim) if not exists(to_logits) else to_logits
 
         # memory tokens (like [cls]) from Memory Transformers paper
 
@@ -2377,13 +2397,12 @@ class TransformerWrapper(Module):
         self.can_cache_kv_outside_max_seq_len = no_abs_pos_emb
 
     def init_(self):
+        if hasattr(self.token_emb, 'init_'):
+            self.token_emb.init_()
+
         if self.l2norm_embed:
-            nn.init.normal_(self.token_emb.emb.weight, std = 1e-5)
             if not isinstance(self.pos_emb, always):
                 nn.init.normal_(self.pos_emb.emb.weight, std = 1e-5)
-            return
-
-        nn.init.kaiming_normal_(self.token_emb.emb.weight)
 
     def forward(
         self,
@@ -2406,7 +2425,9 @@ class TransformerWrapper(Module):
         attn_z_loss_weight = 1e-4,
         seq_start_pos = None,
         cache: LayerIntermediates | None = None,
-        **kwargs
+        token_emb_kwargs = dict(),
+        to_logits_kwargs = dict(),
+        **kwargs,
     ):
         b, n, device, num_mems, has_memory_tokens, emb_frac_gradient, orig_mask = x.shape[0], x.shape[1], x.device, self.num_memory_tokens, self.num_memory_tokens > 0, self.emb_frac_gradient, mask
 
@@ -2417,7 +2438,7 @@ class TransformerWrapper(Module):
 
         external_pos_emb = exists(pos) and pos.dtype != torch.long
         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
+        x = self.token_emb(x, **token_emb_kwargs) + pos_emb
 
         # add additional embeddings
 
@@ -2572,9 +2593,9 @@ class TransformerWrapper(Module):
 
         if not return_embeddings:
             if self.has_multiple_heads:
-                logits = tuple(fn(x) for fn in self.to_logits)
+                logits = tuple(fn(x, **to_logits_kwargs) for fn in self.to_logits)
             else:
-                logits = self.to_logits(x)
+                logits = self.to_logits(x, **to_logits_kwargs)
 
         # maybe sig softmax
 

{x_transformers-1.42.6.dist-info → x_transformers-1.42.8.dist-info}/METADATA

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

{x_transformers-1.42.6.dist-info → x_transformers-1.42.8.dist-info}/RECORD

@@ -6,11 +6,11 @@ x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
-x_transformers/x_transformers.py,sha256=cPsSl1s14_c9fMdn9cZwe6Eg3aDbcRyCTsoXUJusWUg,92706
+x_transformers/x_transformers.py,sha256=275B_yDHePxUvlLcMNgnCUmZ1qZEkwBrpk6IA8n-pnY,93550
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-1.42.6.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.42.6.dist-info/METADATA,sha256=OANeMK9I504gC7iErAdYMTGBUEl6FOcEwm97o4OyC1k,689
-x_transformers-1.42.6.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-x_transformers-1.42.6.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.42.6.dist-info/RECORD,,
+x_transformers-1.42.8.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.42.8.dist-info/METADATA,sha256=1d2BVA6iHKpT4UzbYxw16ijAFGJT-u29zTnYtV6Lp3w,689
+x_transformers-1.42.8.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+x_transformers-1.42.8.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.42.8.dist-info/RECORD,,