x-transformers 1.21.2__py3-none-any.whl → 1.21.4__py3-none-any.whl

x_transformers/x_transformers.py

@@ -11,7 +11,7 @@ from collections import namedtuple
 from dataclasses import dataclass
 from typing import List, Callable, Optional
 
-from einops import rearrange, repeat, reduce
+from einops import rearrange, repeat, reduce, pack, unpack
 from einops.layers.torch import Rearrange
 
 from x_transformers.attend import Attend, Intermediates, CascadingHeads
@@ -27,6 +27,7 @@ class LayerIntermediates:
     attn_intermediates: Optional[List[Intermediates]] = None
     layer_hiddens: Optional[List[Tensor]] = None
     attn_z_loss: Optional[Tensor] = None
+    mems: Optional[Tensor] = None
 
 # helpers
 
@@ -778,8 +779,8 @@ class Attention(nn.Module):
         r_input = x
 
         if exists(mem):
-            k_input = torch.cat((mem, k_input), dim = -2)
-            v_input = torch.cat((mem, v_input), dim = -2)
+            k_input, mem_packed_shape = pack([mem, k_input], 'b * d')
+            v_input, _ = pack([mem, v_input], 'b * d')
 
         q = self.to_q(q_input)
         k = self.to_k(k_input)
@@ -792,11 +793,21 @@ class Attention(nn.Module):
 
         if exists(cache) and not has_context:
             ck, cv = cache.cached_kv
+
+            if exists(mem):
+                mk, k = unpack(k, mem_packed_shape, 'b h * d')
+                mv, v = unpack(v, mem_packed_shape, 'b h * d')
+
             k = torch.cat((ck, k), dim = -2)
             v = torch.cat((cv, v), dim = -2)
 
+            if exists(mem):
+                k = torch.cat((mk, k), dim = -2)
+                v = torch.cat((mv, v), dim = -2)
+
         if return_intermediates:
-            cached_kv = (k, v)
+            mem_len = mem.shape[-2] if exists(mem) else 0
+            cached_kv = (k[..., mem_len:, :], v[..., mem_len:, :])
 
         if self.qk_norm:
             qk_l2norm = partial(l2norm, groups = self.qk_norm_groups)
@@ -941,6 +952,8 @@ class AttentionLayers(nn.Module):
         custom_layers = None,
         sandwich_coef = None,
         par_ratio = None,
+        weight_tie_layers = False,   # Albert - https://arxiv.org/abs/1909.11942
+        layers_execute_order = None, # generalizes weight tying, can do arbitrary layer execution orders
         residual_attn = False,
         cross_residual_attn = False,
         macaron = False,
@@ -1046,6 +1059,15 @@ class AttentionLayers(nn.Module):
             attn_kwargs = {**attn_kwargs, 'zero_init_output':  True}
             ff_kwargs = {**ff_kwargs, 'zero_init_output':  True}
 
+        # setup weight tying, which is a special case of `layer_execute_order`
+
+        assert not (weight_tie_layers and any([*map(exists, (custom_layers, par_ratio, sandwich_coef))]))
+
+        if weight_tie_layers:
+            assert not exists(layers_execute_order)
+            layers_execute_order = tuple(range(len(default_block))) * depth
+            depth = 1
+
         # calculate layer block order
 
         if exists(custom_layers):
@@ -1068,6 +1090,10 @@ class AttentionLayers(nn.Module):
             layer_types = default_block * depth
 
         self.layer_types = layer_types
+        self.layers_execute_order = default(layers_execute_order, tuple(range(len(layer_types))))
+
+        assert all([i < len(self.layer_types) for i in self.layers_execute_order])
+
         self.num_attn_layers = len(list(filter(equals('a'), layer_types)))
 
         # stochastic depth
@@ -1176,7 +1202,19 @@ class AttentionLayers(nn.Module):
 
         outer_residual = x * self.resi_dual_scale
 
-        for ind, (layer_type, (norm, block, residual_fn), layer_dropout) in enumerate(zip(self.layer_types, self.layers, self.layer_dropouts)):
+        # get layers to be executed
+
+        layer_variables = (
+            self.layer_types,
+            self.layers,
+            self.layer_dropouts
+        )
+
+        layer_variables = tuple(tuple(layer_variable[i] for i in self.layers_execute_order) for layer_variable in layer_variables)
+
+        # go through the attention and feedforward layers
+
+        for ind, (layer_type, (norm, block, residual_fn), layer_dropout) in enumerate(zip(*layer_variables)):
             is_last = ind == (len(self.layers) - 1)
 
             if self.training and layer_dropout > 0. and random() < layer_dropout:
@@ -1475,14 +1513,18 @@ class TransformerWrapper(nn.Module):
             intermediates.attn_z_loss = calc_z_loss(pre_softmax_attns, weight = attn_z_loss_weight)
             return_intermediates = True
 
-        if return_intermediates:
-            return out, intermediates
-
         if return_mems:
             hiddens = intermediates.hiddens
             new_mems = list(map(lambda pair: torch.cat(pair, dim = -2), zip(mems, hiddens))) if exists(mems) else hiddens
             new_mems = list(map(lambda t: t[..., -self.max_mem_len:, :].detach(), new_mems))
-            return out, new_mems
+
+            if not return_intermediates:
+                return out, new_mems
+
+            intermediates.mems = new_mems
+
+        if return_intermediates:
+            return out, intermediates
 
         if return_attn:
             attn_maps = list(map(lambda t: t.post_softmax_attn, intermediates.attn_intermediates))

x_transformers/xl_autoregressive_wrapper.py

@@ -67,6 +67,7 @@ class XLAutoregressiveWrapper(nn.Module):
         curr_pos = len(all_leading_tokens) * max_seq_len
         curr_mems = mems
 
+        cache = None
         out = start_tokens
 
         for _ in range(seq_len):
@@ -75,13 +76,17 @@ class XLAutoregressiveWrapper(nn.Module):
 
             x = out[:, curr_pos:]
 
-            logits, mems = self.net(
+            logits, cache = self.net(
                 x,
                 mems = curr_mems,
+                cache = cache,
                 return_mems = True,
+                return_intermediates = True,
                 **kwargs
             )
 
+            mems = cache.mems
+
             logits = logits[:, -1]
             filtered_logits = filter_logits_fn(logits, thres = filter_thres)
             probs = F.softmax(filtered_logits / temperature, dim=-1)

{x_transformers-1.21.2.dist-info → x_transformers-1.21.4.dist-info}/METADATA

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

{x_transformers-1.21.2.dist-info → x_transformers-1.21.4.dist-info}/RECORD

@@ -3,10 +3,10 @@ x_transformers/attend.py,sha256=_LepMOwph_o3jio6tur8LEUBPM-2YIn7NpuBmqhU47E,1238
 x_transformers/autoregressive_wrapper.py,sha256=djA4nfE6_92SAzX1JI0KaC7krdLz1mvnZlaVOaerHDg,5372
 x_transformers/continuous_autoregressive_wrapper.py,sha256=pTiDqu6JRUlnQJQp_xHATYHy0lgSd6ERLqyiFO3pC-4,1575
 x_transformers/nonautoregressive_wrapper.py,sha256=AQLE4rA_Kh8VNoe9OzpwyeWson34sRkhks4dn4seNjI,10414
-x_transformers/x_transformers.py,sha256=Sw5FJlyDQ0fEINLw3G9aHJlLy3EBir0NphBVZ_gLiuA,56508
-x_transformers/xl_autoregressive_wrapper.py,sha256=-CAYjTtqrks8ZTxjYm2stOelZpU4MbZIvLjUxWO0P9Y,3988
-x_transformers-1.21.2.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.21.2.dist-info/METADATA,sha256=q47HwjcwIc31Liw_UWeHQKhjRZaWjrSLQo5ryHgCfkQ,661
-x_transformers-1.21.2.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-x_transformers-1.21.2.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.21.2.dist-info/RECORD,,
+x_transformers/x_transformers.py,sha256=vss6ISABCV74wgEapIt8nPK50j9-QU54hugRiaBh-sw,58088
+x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
+x_transformers-1.21.4.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.21.4.dist-info/METADATA,sha256=HzzezgisQhEH2H6D2tI-JDqNghpuUi6pDlHg0AI976U,661
+x_transformers-1.21.4.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
+x_transformers-1.21.4.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.21.4.dist-info/RECORD,,