x-transformers 1.35.2__py3-none-any.whl → 1.36.0__py3-none-any.whl

x_transformers/x_transformers.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import math
-from random import random
+from random import random, randrange
 from packaging import version
 
 import torch
@@ -12,6 +12,7 @@ from torch.amp import autocast
 
 from functools import partial, wraps
 from collections import namedtuple
+from contextlib import nullcontext
 from dataclasses import dataclass
 from typing import List, Dict, Tuple, Callable
 
@@ -920,6 +921,7 @@ class Attention(Module):
         kv_heads = None,
         shared_kv = False,
         value_dim_head = None,
+        dim_out = None,
         tensor_product = False,      # https://arxiv.org/abs/2208.06061
         add_zero_kv = False,         # same as add_zero_attn in pytorch
         rotary_embed_values = False,
@@ -1057,7 +1059,11 @@ class Attention(Module):
         # attention on attention
 
         self.attn_on_attn = on_attn
-        self.to_out = nn.Sequential(nn.Linear(out_dim, dim * 2, bias = False), nn.GLU()) if on_attn else nn.Linear(out_dim, dim, bias = False)
+
+        # output dimension by default same as input, but can be overridden
+
+        dim_out = default(dim_out, dim)
+        self.to_out = nn.Sequential(nn.Linear(out_dim, dim_out * 2, bias = False), nn.GLU()) if on_attn else nn.Linear(out_dim, dim_out, bias = False)
 
         # whether to rotate positions into values, for absolute positions in addition to relative
 
@@ -1982,7 +1988,9 @@ class TransformerWrapper(Module):
         use_abs_pos_emb = True,
         scaled_sinu_pos_emb = False,
         l2norm_embed = False,
-        emb_frac_gradient = 1., # GLM-130B and Cogview successfully used this, set at 0.1
+        recycling = False,            # from Jumper et al. - Alphafold2
+        train_max_recycle_steps = 4,  # saw a benefit for language modeling up to 3 recycling steps, so let's default this to 4
+        emb_frac_gradient = 1.,       # GLM-130B and Cogview successfully used this, set at 0.1
         attn_z_loss_weight = 1e-4,
         average_pool_embed = False,
         use_cls_token = False,
@@ -2039,6 +2047,13 @@ class TransformerWrapper(Module):
 
         assert at_most_one_of(average_pool_embed, use_cls_token)
 
+        # maybe recycling
+
+        self.recycling = recycling
+        self.recycled_proj = nn.Linear(dim, dim, bias = False) if recycling else None
+
+        self.train_max_recycle_steps = train_max_recycle_steps
+
         # classic cls token from the bert days
 
         self.cls_token = None
@@ -2082,7 +2097,7 @@ class TransformerWrapper(Module):
 
         # whether can do cached kv decoding
 
-        self.can_cache_kv = self.num_memory_tokens == 0
+        self.can_cache_kv = self.num_memory_tokens == 0 and not recycling
         self.can_cache_kv_outside_max_seq_len = no_abs_pos_emb
 
     def init_(self):
@@ -2105,6 +2120,7 @@ class TransformerWrapper(Module):
         return_attn = False,
         mems = None,
         mem_masks = None,
+        recycle_steps = None,
         pos = None,
         prepend_embeds = None,
         prepend_mask = None,
@@ -2210,11 +2226,37 @@ class TransformerWrapper(Module):
             if exists(mem_every):
                 x = rearrange(x, '(b n) m d -> b (n m) d', b = b)
 
+        # handle maybe shifting of memories
+
         if self.shift_mem_down and exists(mems):
             mems_l, mems_r = mems[:self.shift_mem_down], mems[self.shift_mem_down:]
             mems = [*mems_r, *mems_l]
 
-        x, intermediates = self.attn_layers(x, mask = mask, mems = mems, mem_masks = mem_masks, cache = cache, return_hiddens = True, seq_start_pos = seq_start_pos, **kwargs)
+        # attention layers
+
+        if not self.recycling:
+            # regular
+
+            attended, intermediates = self.attn_layers(x, mask = mask, mems = mems, mem_masks = mem_masks, cache = cache, return_hiddens = True, seq_start_pos = seq_start_pos, **kwargs)
+
+        else:
+            # recycling
+
+            recycle_steps = default(recycle_steps, (randrange(self.train_max_recycle_steps) + 1) if self.training else None)
+            assert exists(recycle_steps) and recycle_steps > 0, '`recycle_steps` must be provided on forward if recycling is turned on and not training'
+
+            for i in range(recycle_steps):
+                first_step = i == 0
+                last_step = i == (recycle_steps - 1)
+
+                context = nullcontext if last_step else torch.no_grad
+
+                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)
+
+        x = attended
 
         # handle memories post-attention
 

{x_transformers-1.35.2.dist-info → x_transformers-1.36.0.dist-info}/METADATA

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

{x_transformers-1.35.2.dist-info → x_transformers-1.36.0.dist-info}/RECORD

@@ -5,11 +5,11 @@ x_transformers/continuous.py,sha256=cIVEdhfei258__ziV7kQBrJMxCel54bExBTDrO9rfCI,
 x_transformers/dpo.py,sha256=LjvWgCkqTl-UuehrzQ8nkX5guLr4whYwsmm7SKSwdls,3450
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/nonautoregressive_wrapper.py,sha256=ys_p8obc7lTeeodCqvkRKxOXQ1C9T3j5Jwr-JbVgnXk,10432
-x_transformers/x_transformers.py,sha256=mpA9hriHqCXLckdlVo8sxzXT6sjxwsY6AaKoP-Rpw3c,80631
+x_transformers/x_transformers.py,sha256=iib15Squ9VE7tLpb4Z4_Hq_hi7dZhPNR_xPtC9BzMrE,82321
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-x_transformers-1.35.2.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.35.2.dist-info/METADATA,sha256=4UuWPhkRRayYadZ8kwaHyqEGhHurnWJGRbPTzDMdEZo,661
-x_transformers-1.35.2.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-x_transformers-1.35.2.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.35.2.dist-info/RECORD,,
+x_transformers-1.36.0.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.36.0.dist-info/METADATA,sha256=YKcnT5T0UkZxwpP72cPfx9RN0SVoBYy0e6Xo581YCE0,661
+x_transformers-1.36.0.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+x_transformers-1.36.0.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.36.0.dist-info/RECORD,,