x-transformers 2.3.18__py3-none-any.whl → 2.3.20__py3-none-any.whl

x_transformers/attend.py

@@ -25,6 +25,7 @@ class Intermediates:
     values:             Tensor | None = None
     cached_kv:          Tuple[Tensor, Tensor] | None = None
     layer_type:         str | None = None
+    hybrid_hidden:      Tensor | None = None
 
     def to_tuple(self):
         return (self.qk_similarities, self.pre_softmax_attn, self.post_softmax_attn)

x_transformers/continuous.py

@@ -32,6 +32,15 @@ def default(val, d):
         return val
     return d() if not isinstance(d, Module) and callable(d) else d
 
+def sample_from_mean_variance(
+    mean,
+    variance,
+    eps = 1e-5,
+    temperature = 1.
+):
+    std = variance.clamp(min = eps).sqrt()
+    return torch.normal(mean, std * temperature)
+
 def masked_mean(t, mask):
     t = einx.where('b n, b n d, -> b n d', mask, t, 0.)
 
@@ -274,9 +283,7 @@ class ContinuousAutoregressiveWrapper(Module):
 
             if self.probabilistic:
                 mean, var = last_output
-                stddev = var.clamp(min = 1e-5).sqrt()
-
-                last_output = torch.normal(mean, stddev * temperature)
+                last_output = sample_from_mean_variance(mean, var, temperature = temperature)
 
             out = cat((out, last_output), dim = -2)
 
@@ -372,8 +379,7 @@ class ContinuousAutoregressiveWrapper(Module):
 
             if self.probabilistic:
                 mean, var = last_pred
-                std = var.clamp(min = 1e-5).sqrt()
-                inp = torch.normal(mean, std)
+                inp = sample_from_mean_variance(mean, var)
             else:
                 inp = last_pred
 

x_transformers/x_transformers.py

@@ -49,6 +49,7 @@ class LayerIntermediates:
     mems:               Tensor | None = None
     memory_tokens:      Tensor | None = None
     logit_entropies:    Tensor | None = None
+    cache_length:       int = 0
 
 LinearNoBias = partial(nn.Linear, bias = False)
 
@@ -282,12 +283,18 @@ class AbsolutePositionalEmbedding(Module):
         self.l2norm_embed = l2norm_embed
         self.emb = nn.Embedding(max_seq_len, dim)
 
-    def forward(self, x, pos = None, seq_start_pos = None):
+    def forward(
+        self,
+        x,
+        pos = None,
+        seq_start_pos = None,
+        offset = 0
+    ):
         seq_len, device = x.shape[1], x.device
         assert seq_len <= self.max_seq_len, f'you are passing in a sequence length of {seq_len} but your absolute positional embedding has a max sequence length of {self.max_seq_len}'
 
         if not exists(pos):
-            pos = arange(seq_len, device = device)
+            pos = arange(seq_len, device = device) + offset
 
         if exists(seq_start_pos):
             pos = (pos - seq_start_pos[..., None]).clamp(min = 0)
@@ -307,11 +314,17 @@ class ScaledSinusoidalEmbedding(Module):
         inv_freq = theta ** -freq_seq
         self.register_buffer('inv_freq', inv_freq, persistent = False)
 
-    def forward(self, x, pos = None, seq_start_pos = None):
+    def forward(
+        self,
+        x,
+        pos = None,
+        seq_start_pos = None,
+        offset = 0
+    ):
         seq_len, device = x.shape[1], x.device
 
         if not exists(pos):
-            pos = arange(seq_len, device = device)
+            pos = arange(seq_len, device = device) + offset
 
         if exists(seq_start_pos):
             pos = pos - seq_start_pos[..., None]
@@ -676,7 +689,7 @@ class RotaryEmbedding(Module):
         return self.forward(t)
 
     @autocast('cuda', enabled = False)
-    def forward(self, t):
+    def forward(self, t, offset = 0):
         max_pos = t.max() + 1
 
         if t.ndim == 1:
@@ -1079,10 +1092,11 @@ class FoldAxially(Module):
     def forward(
         self,
         x,
+        *args,
         **kwargs
     ):
         if self.axial_dim == 1:
-            return self.fn(x, **kwargs)
+            return self.fn(x, *args, **kwargs)
 
         seq_len, axial_dim = x.shape[1], self.axial_dim
 
@@ -1091,7 +1105,7 @@ class FoldAxially(Module):
 
         x = rearrange(x, 'b (n axial_dim) ... -> (b axial_dim) n ...', axial_dim = axial_dim)
 
-        out = self.fn(x, **kwargs)
+        out = self.fn(x, *args, **kwargs)
 
         (out, *rest_out), tree_spec = tree_flatten(out)
 
@@ -1857,9 +1871,17 @@ class Attention(Module):
             if not self.causal and exists(self.hybrid_mask_kwarg):
                 hybrid_forward_kwargs = {self.hybrid_mask_kwarg: mask}
 
+            # handle maybe hybrid cache
+
+            hybrid_forward_args = ()
+
+            if exists(cache) and exists(cache.hybrid_hidden):
+                hybrid_hiddens = cache.hybrid_hidden
+                hybrid_forward_args = (hybrid_hiddens,)
+
             # hybrid forward
 
-            hybrid_outputs = self.hybrid_module(x, **hybrid_forward_kwargs)
+            hybrid_outputs = self.hybrid_module(x, *hybrid_forward_args, **hybrid_forward_kwargs)
 
             # handle hybrid out
 
@@ -1870,6 +1892,10 @@ class Attention(Module):
             if hybrid_out.ndim == 3:
                 hybrid_out = rearrange(hybrid_out, 'b n (h d) -> b h n d', h = h)
 
+            if len(rest_hybrid_outs) > 0:
+                hybrid_hidden = first(rest_hybrid_outs)
+                intermediates.hybrid_hidden = hybrid_hidden
+
             out_norm, hybrid_out_norm = self.hybrid_norms
 
             out = out_norm(out)
@@ -2360,7 +2386,9 @@ class AttentionLayers(Module):
         mems = None,
         mem_masks = None,
         seq_start_pos: Tensor | None = None,
+        seq_pos_offset: int = 0,
         cache: LayerIntermediates | None = None,
+        input_not_include_cache = False,
         cache_age = 1,
         return_hiddens = False,
         rotary_pos_emb = None,
@@ -2434,7 +2462,7 @@ class AttentionLayers(Module):
                 mem_len = maybe_mem.shape[1] if exists(maybe_mem) else 0
 
                 if not exists(pos):
-                    pos = arange(x.shape[1] + mem_len, device = x.device) - mem_len
+                    pos = arange(x.shape[1] + mem_len + seq_pos_offset, device = x.device) - mem_len
 
                 rotary_pos_emb = self.rotary_pos_emb(pos)
 
@@ -2451,11 +2479,15 @@ class AttentionLayers(Module):
 
         # assume cached key / values
 
+        prev_cache_length = 0
+
         attn_cache = []
 
         if exists(cache):
             assert self.causal and not any([*map(exists, (mask, attn_mask))])
 
+            prev_cache_length = cache.cache_length
+
             if exists(context):
                 context = context[:, :0]
 
@@ -2469,6 +2501,8 @@ class AttentionLayers(Module):
 
             attn_cache = cache.attn_intermediates
 
+        next_cache_length = x.shape[1]
+
         iter_attn_cache = iter(attn_cache)
 
         # handle deep embeds if needed
@@ -2655,6 +2689,7 @@ class AttentionLayers(Module):
             last_hidden = x,
             attn_intermediates = intermediates,
             layer_hiddens = layer_hiddens,
+            cache_length = next_cache_length + prev_cache_length
         )
 
         return x, intermediates
@@ -2989,6 +3024,7 @@ class TransformerWrapper(Module):
         attn_z_loss_weight = 1e-4,
         seq_start_pos = None,
         cache: LayerIntermediates | None = None,
+        input_not_include_cache = False,
         token_emb_kwargs = dict(),
         to_logits_kwargs = dict(),
         **kwargs,
@@ -3007,10 +3043,17 @@ class TransformerWrapper(Module):
         return_hiddens = return_mems | return_attn | return_intermediates | return_attn_z_loss | return_embeddings_and_intermediates
         return_embeddings = return_embeddings | (not exists(self.to_logits)) | return_embeddings_and_intermediates
 
+        # take care of position embedding offsets in the presence of cache and sequence is less than cache length (not full sequence)
+
+        seq_pos_offset = 0
+
+        if exists(cache) and input_not_include_cache:
+            seq_pos_offset = cache.cache_length
+
         # absolute positional embedding
 
         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
+        pos_emb = self.pos_emb(x, pos = pos, seq_start_pos = seq_start_pos, offset = seq_pos_offset) if not external_pos_emb else pos
         x = self.token_emb(x, **token_emb_kwargs) + pos_emb
 
         # add additional embeddings
@@ -3109,6 +3152,15 @@ class TransformerWrapper(Module):
             mems_l, mems_r = mems[:self.shift_mem_down], mems[self.shift_mem_down:]
             mems = [*mems_r, *mems_l]
 
+        # attn layers kwargs
+
+        kwargs = dict(
+            **kwargs,
+            seq_pos_offset = seq_pos_offset,
+            seq_start_pos = seq_start_pos,
+            input_not_include_cache = input_not_include_cache
+        )
+
         # attention layers
 
         if not self.recycling:
@@ -3116,7 +3168,7 @@ class TransformerWrapper(Module):
 
             # regular
 
-            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)
+            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, **kwargs)
 
         else:
             # recycling
@@ -3133,7 +3185,7 @@ class TransformerWrapper(Module):
                 with context():
                     maybe_recycled = self.recycled_proj(attended.detach()) if not first_step else 0.
 
-                    attended, intermediates = self.attn_layers(x + maybe_recycled, mask = mask, mems = mems, mem_masks = mem_masks, cache = cache, return_hiddens = True, seq_start_pos = seq_start_pos, **kwargs)
+                    attended, intermediates = self.attn_layers(x + maybe_recycled, mask = mask, mems = mems, mem_masks = mem_masks, cache = cache, return_hiddens = True, **kwargs)
 
         x = attended
 

{x_transformers-2.3.18.dist-info → x_transformers-2.3.20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.3.18
+Version: 2.3.20
 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.18.dist-info → x_transformers-2.3.20.dist-info}/RECORD

@@ -1,17 +1,17 @@
 x_transformers/__init__.py,sha256=h3I2ejobgEdy8H7NgV-rP8UaBCnd16-MysvDXH9GMEA,985
-x_transformers/attend.py,sha256=xFsBtl7h7_qebPh7kE81BpmCWAjCgFpB9i_IHu_91es,17288
+x_transformers/attend.py,sha256=fXMuwHuBAFB4f4_U6j5_uVeK7N4cV0PDd6UTqtkjKKM,17333
 x_transformers/autoregressive_wrapper.py,sha256=LW1gr3cFONDEPA_HHhaTE7mk-JWbaINuB1fc_DfbCqw,10791
 x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTNY__GGZEKU,13374
-x_transformers/continuous.py,sha256=uV2hLQOckeRsybqJy-0F8RhAyMPJlkVHmA7QqUJHG4g,12433
+x_transformers/continuous.py,sha256=CHta8vizKl85n220fv5278fwjSU-vrN_FBy-m831_go,12551
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/entropy_based_tokenizer.py,sha256=F2lO8-v3aLIcVDVNhu7RR-UtRdlmaaYJzBK9m7OnLE8,5018
 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=ZfOXrZSiy2jlZ8wVmDdMTLW4hAY_qfmPQHW9t2ABxbo,114097
+x_transformers/x_transformers.py,sha256=l2p-r0iJNlYHUB3vM4lb6ptzNCx9HgA7UfgieEcQT6w,115521
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.3.18.dist-info/METADATA,sha256=RKXNlO50fifu1Nas38iZRn6IJVDkv4Cen94XYVJlWg0,89897
-x_transformers-2.3.18.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.3.18.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.3.18.dist-info/RECORD,,
+x_transformers-2.3.20.dist-info/METADATA,sha256=ygWyfnlIh2Mw6bd12gJjjZJyM9vfnXmvvOLyrd2El2k,89897
+x_transformers-2.3.20.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.3.20.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.3.20.dist-info/RECORD,,