x-transformers 2.3.7__py3-none-any.whl → 2.3.9__py3-none-any.whl

x_transformers/x_transformers.py

@@ -1208,8 +1208,6 @@ class FeedForward(Module):
         sublayer_dropout = 0.,
         no_bias = False,
         zero_init_output = False,
-        deep_embed_hiddens = False,
-        deep_embed_num_tokens = None,
     ):
         super().__init__()
         inner_dim = int(dim * mult)
@@ -1242,17 +1240,6 @@ class FeedForward(Module):
             nn.Dropout(sublayer_dropout) if sublayer_dropout > 0. else None
         )
 
-        # deep embed
-
-        # credit goes to Braden Koszarsky for first devising value embeddings in nanogpt-speedrun project
-        # then Bo Peng for coming up with this alternate design in feedforward for RWKV 8
-        # improvements were clearest to me (on my toy setup) with multiplying on output of feedforward, will try with attention at future date
-
-        self.deep_embed = None
-        if deep_embed_hiddens:
-            assert exists(deep_embed_num_tokens)
-            self.deep_embed = nn.Parameter(torch.ones(deep_embed_num_tokens, dim_out))
-
         # init last linear layer to 0
 
         if zero_init_output:
@@ -1261,12 +1248,11 @@ class FeedForward(Module):
     def forward(
         self,
         x,
-        deep_embed_ids = None
+        deep_embed = None
     ):
         out = self.ff(x)
 
-        if exists(deep_embed_ids) and exists(self.deep_embed):
-            deep_embed = self.deep_embed[deep_embed_ids]
+        if exists(deep_embed):
             out = out * deep_embed
 
         return out
@@ -2380,7 +2366,7 @@ class AttentionLayers(Module):
         pos = None,
         context_pos = None,
         attn_bias = None,
-        deep_embed_ids = None,
+        deep_embeds_and_ids: tuple[nn.Parameter, Tensor] | None = None,
         condition = None,
         in_attn_cond = None, # https://arxiv.org/abs/2105.04090
         layers_execute_order: tuple[int, ...] | None = None
@@ -2475,13 +2461,26 @@ class AttentionLayers(Module):
             if cache_age > 0:
                 x = x[:, -cache_age:] # for spec decoding, may be greater than 1
 
-                if exists(deep_embed_ids):
-                    deep_embed_ids = deep_embed_ids[:, -cache_age:]
+                if exists(deep_embeds_and_ids):
+                    deep_embeds, token_ids = deep_embeds_and_ids
+                    token_ids = token_ids[:, -cache_age:]
+                    deep_embeds_and_ids = (deep_embeds, token_ids)
 
             attn_cache = cache.attn_intermediates
 
         iter_attn_cache = iter(attn_cache)
 
+        # handle deep embeds if needed
+
+        deep_embeds = []
+
+        if exists(deep_embeds_and_ids):
+            deep_embeds, token_ids = deep_embeds_and_ids
+            deep_embeds_across_depth = deep_embeds[token_ids]
+            deep_embeds = rearrange(deep_embeds_across_depth, 'b n l d -> l b n d')
+
+        deep_embeds_iter = iter(deep_embeds)
+
         # setup multistreams if needed
 
         streams = self.num_residual_streams
@@ -2602,7 +2601,7 @@ class AttentionLayers(Module):
             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, **cross_attn_rotary_pos_emb, return_intermediates = True)
             elif layer_type == 'f':
-                out = block(x, deep_embed_ids = deep_embed_ids)
+                out = block(x, deep_embed = next(deep_embeds_iter, None))
 
             # store first self or cross attention intermediate for value residual
 
@@ -2818,11 +2817,14 @@ class TransformerWrapper(Module):
         mixture_of_softmax = False,
         mixture_of_softmax_k = 4,
         sigsoftmax_logits = False,
+        ff_deep_embed = False,
         to_logits: Module | None = None,
     ):
         super().__init__()
 
         dim = attn_layers.dim
+        depth = attn_layers.depth
+
         emb_dim = default(emb_dim, dim)
         self.emb_dim = emb_dim
         self.num_tokens = num_tokens
@@ -2855,6 +2857,16 @@ class TransformerWrapper(Module):
         if len(embed_num_tokens) > 0:
             self.embeds = ModuleDict({f'{name}_embed': nn.Embedding(num_tokens, emb_dim) for name, num_tokens in embed_num_tokens.items()})
 
+        # deep embed
+
+        # credit goes to Braden Koszarsky for first devising value embeddings in nanogpt-speedrun project
+        # then Bo Peng for coming up with this alternate design in feedforward for RWKV 8
+        # improvements were clearest to me (on my toy setup) with multiplying on output of feedforward, will try with attention at future date
+
+        self.ff_deep_embed = None
+        if ff_deep_embed:
+            self.ff_deep_embed = nn.Parameter(torch.ones(num_tokens, depth, dim))
+
         # fraction of the gradient that should go to the embedding, https://arxiv.org/abs/2105.13290
 
         self.emb_frac_gradient = emb_frac_gradient
@@ -3050,6 +3062,13 @@ class TransformerWrapper(Module):
 
         x = self.project_emb(x)
 
+        # maybe deep embeds
+
+        deep_embed_and_ids = None
+
+        if exists(self.ff_deep_embed):
+            deep_embed_and_ids = (self.ff_deep_embed, token_ids)
+
         # maybe cls token
 
         if exists(self.cls_token):
@@ -3096,7 +3115,7 @@ class TransformerWrapper(Module):
 
             # regular
 
-            attended, intermediates = self.attn_layers(x, mask = mask, mems = mems, mem_masks = mem_masks, cache = cache, deep_embed_ids = token_ids, return_hiddens = True, seq_start_pos = seq_start_pos, **kwargs)
+            attended, intermediates = self.attn_layers(x, mask = mask, mems = mems, mem_masks = mem_masks, cache = cache, deep_embeds_and_ids = deep_embed_and_ids, return_hiddens = True, seq_start_pos = seq_start_pos, **kwargs)
 
         else:
             # recycling

{x_transformers-2.3.7.dist-info → x_transformers-2.3.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.3.7
+Version: 2.3.9
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers

{x_transformers-2.3.7.dist-info → x_transformers-2.3.9.dist-info}/RECORD

@@ -8,10 +8,10 @@ x_transformers/entropy_based_tokenizer.py,sha256=F2lO8-v3aLIcVDVNhu7RR-UtRdlmaaY
 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=bwuZlvse3kYpD0EoHM9gWyi5IpXkF-jsNgQqJGjfRzs,113501
+x_transformers/x_transformers.py,sha256=Wo5hauzdn4Q9PUVjBqQo-1vCq08BT2jYUDbq3r2a5Go,114061
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-2.3.7.dist-info/METADATA,sha256=U80x0At1b-5MP3co5wM7CdL7zmtTHWjbYzXn5ypEBoU,89021
-x_transformers-2.3.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.3.7.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.3.7.dist-info/RECORD,,
+x_transformers-2.3.9.dist-info/METADATA,sha256=M0rUg95el9swG5VFWas-5YjahLSgytT65W9L3Ne7BJM,89021
+x_transformers-2.3.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.3.9.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.3.9.dist-info/RECORD,,