x-transformers 2.3.21__py3-none-any.whl → 2.3.23__py3-none-any.whl

x_transformers/attend.py

@@ -276,21 +276,22 @@ class Attend(Module):
 
         # torch 2.3 uses new backend and context manager
 
-        if torch_version >= version.parse('2.3'):
-            from torch.nn.attention import SDPBackend
-
-            str_to_backend = dict(
-                enable_flash = SDPBackend.FLASH_ATTENTION,
-                enable_mem_efficient = SDPBackend.EFFICIENT_ATTENTION,
-                enable_math = SDPBackend.MATH,
-                enable_cudnn = SDPBackend.CUDNN_ATTENTION
-            )
+        if self.flash:
+            if torch_version >= version.parse('2.3'):
+                from torch.nn.attention import SDPBackend
 
-            sdpa_backends = [str_to_backend[enable_str] for enable_str, enable in sdp_kwargs.items() if enable]
+                str_to_backend = dict(
+                    enable_flash = SDPBackend.FLASH_ATTENTION,
+                    enable_mem_efficient = SDPBackend.EFFICIENT_ATTENTION,
+                    enable_math = SDPBackend.MATH,
+                    enable_cudnn = SDPBackend.CUDNN_ATTENTION
+                )
 
-            self.sdp_context_manager = partial(torch.nn.attention.sdpa_kernel, sdpa_backends)
-        else:
-            self.sdp_context_manager = partial(torch.backends.cuda.sdp_kernel, **sdp_kwargs)
+                sdpa_backends = [str_to_backend[enable_str] for enable_str, enable in sdp_kwargs.items() if enable]
+
+                self.sdp_context_manager = partial(torch.nn.attention.sdpa_kernel, sdpa_backends)
+            else:
+                self.sdp_context_manager = partial(torch.backends.cuda.sdp_kernel, **sdp_kwargs)
 
     def flash_attn(
         self,

x_transformers/autoregressive_wrapper.py

@@ -140,7 +140,8 @@ class AutoregressiveWrapper(Module):
         ignore_index = -100,
         pad_value = 0,
         mask_prob = 0.,
-        add_attn_z_loss = False
+        add_attn_z_loss = False,
+        next_embed_loss_weight = 0.1
     ):
         super().__init__()
         self.pad_value = pad_value
@@ -156,6 +157,10 @@ class AutoregressiveWrapper(Module):
         # whether to add router z-loss
         self.add_attn_z_loss = add_attn_z_loss
 
+        # whether to add a continuous loss
+        self.add_continuous_pred_head = net.add_continuous_pred_head
+        self.next_embed_loss_weight = next_embed_loss_weight
+
     @torch.no_grad()
     @eval_decorator
     def generate(
@@ -305,9 +310,9 @@ class AutoregressiveWrapper(Module):
         return out
 
     def forward(self, x, return_outputs = False, **kwargs):
-        seq, ignore_index, add_attn_z_loss = x.shape[1], self.ignore_index, self.add_attn_z_loss
+        seq, ignore_index, add_attn_z_loss, add_next_embed_loss = x.shape[1], self.ignore_index, self.add_attn_z_loss, self.add_continuous_pred_head
 
-        inp, target = x[:, :-1], x[:, 1:]
+        inp, target = x, x[:, 1:]
         inp = torch.where(inp == ignore_index, self.pad_value, inp)
 
         if self.mask_prob > 0.:
@@ -318,15 +323,29 @@ class AutoregressiveWrapper(Module):
             mask = ~torch.zeros_like(inp).scatter(1, indices, 1.).bool()
             kwargs.update(self_attn_kv_mask = mask)
 
-        logits, cache = self.net(
-            inp,
+        out, cache = self.net(
+            x,
             return_intermediates = True,
             return_attn_z_loss = add_attn_z_loss,
+            return_next_embed_pred = add_next_embed_loss,
             **kwargs
         )
 
+        # destruct differently if doing continuous pred
+
+        if add_next_embed_loss:
+            logits, (next_embed_pred, init_embeds) = out
+        else:
+            logits = out
+
+        logits = logits[:, :-1]
+
+        # loss function
+
         loss_fn = F.cross_entropy if not self.net.output_is_log_prob else F.nll_loss
 
+        # cross entropy loss
+
         loss = loss_fn(
             rearrange(logits, 'b n c -> b c n'),
             target,
@@ -336,6 +355,16 @@ class AutoregressiveWrapper(Module):
         if add_attn_z_loss:
             loss = loss + cache.attn_z_loss
 
+        if add_next_embed_loss:
+            mask = inp[:, :-1] != ignore_index
+            embed_pred = next_embed_pred[:, :-1]
+            cont_targets = init_embeds[:, 1:].detach()
+
+            cont_loss = F.l1_loss(embed_pred, cont_targets, reduction = 'none')
+            cont_loss = cont_loss[mask].mean()
+
+            loss = loss + cont_loss * self.next_embed_loss_weight
+
         if not return_outputs:
             return loss
 

x_transformers/x_transformers.py

@@ -2855,6 +2855,7 @@ class TransformerWrapper(Module):
         sigsoftmax_logits = False,
         ff_deep_embed = False,
         to_logits: Module | None = None,
+        add_continuous_pred_head = False
     ):
         super().__init__()
 
@@ -2975,6 +2976,18 @@ class TransformerWrapper(Module):
         else:
             self.to_logits = LinearNoBias(dim, logits_dim) if not exists(to_logits) else to_logits
 
+        # add a head that predicts the embedding of the next step
+
+        self.add_continuous_pred_head = add_continuous_pred_head
+
+        if add_continuous_pred_head:
+
+            self.to_next_embed_pred = nn.Sequential(
+                LinearNoBias(dim, dim),
+                nn.SiLU(),
+                LinearNoBias(dim, dim)
+            )
+
         # memory tokens (like [cls]) from Memory Transformers paper
 
         num_memory_tokens = default(num_memory_tokens, 0)
@@ -3009,6 +3022,7 @@ class TransformerWrapper(Module):
         return_intermediates = False,
         return_embeddings_and_intermediates = False,
         return_logit_entropies = False,
+        return_next_embed_pred = False,
         mask = None,
         return_mems = False,
         return_attn = False,
@@ -3100,6 +3114,10 @@ class TransformerWrapper(Module):
             assert emb_frac_gradient > 0
             x = x * emb_frac_gradient + x.detach() * (1 - emb_frac_gradient)
 
+        # init embed
+
+        init_embed = x
+
         # embedding dropout
 
         x = self.emb_dropout(x)
@@ -3261,6 +3279,14 @@ class TransformerWrapper(Module):
         else:
             out = logits
 
+        # maybe next embed pred
+
+        if return_next_embed_pred:
+            assert self.add_continuous_pred_head
+            next_embed_out = self.to_next_embed_pred(x)
+
+            out = (out, (next_embed_out, init_embed))
+
         # logit entropies
 
         if return_logit_entropies:

{x_transformers-2.3.21.dist-info → x_transformers-2.3.23.dist-info}/METADATA

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

@@ -1,6 +1,6 @@
 x_transformers/__init__.py,sha256=h3I2ejobgEdy8H7NgV-rP8UaBCnd16-MysvDXH9GMEA,985
-x_transformers/attend.py,sha256=fXMuwHuBAFB4f4_U6j5_uVeK7N4cV0PDd6UTqtkjKKM,17333
-x_transformers/autoregressive_wrapper.py,sha256=LW1gr3cFONDEPA_HHhaTE7mk-JWbaINuB1fc_DfbCqw,10791
+x_transformers/attend.py,sha256=Ax34Rw56xXAWfFPqtZ_B8iKEW2EfQdbVoc9uFjfeNjA,17404
+x_transformers/autoregressive_wrapper.py,sha256=BWFaO-3YWzCcEfp-EC1ZkdckqDpPIOQG6_uyyP6AmhM,11753
 x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTNY__GGZEKU,13374
 x_transformers/continuous.py,sha256=hpb1sSbt3k2LNzzjrjSd8F5xOIbKj7IluV9MBEAFLkw,13031
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
@@ -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=l2p-r0iJNlYHUB3vM4lb6ptzNCx9HgA7UfgieEcQT6w,115521
+x_transformers/x_transformers.py,sha256=7phSZvP1_SDRIkVMwVR4cz1dFU2UlR2Wf1HJHEQlcQg,116222
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.3.21.dist-info/METADATA,sha256=530_RGFGFlDyKIV6vMGqjGGw0f3gpArBbwNBHai_LQs,89897
-x_transformers-2.3.21.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.3.21.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.3.21.dist-info/RECORD,,
+x_transformers-2.3.23.dist-info/METADATA,sha256=xRMZP1TSYdcbc0F5GX-WcaHhAbQPdGeFIbjHBZYG9_0,89897
+x_transformers-2.3.23.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.3.23.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.3.23.dist-info/RECORD,,