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

x_transformers/x_transformers.py

@@ -952,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,
@@ -1057,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):
@@ -1079,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
@@ -1187,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:

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

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

{x_transformers-1.21.3.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=NIp-rnED6zAgOUqeOdDh4GzbCz4RbebYXtBVemT7EMc,57031
+x_transformers/x_transformers.py,sha256=vss6ISABCV74wgEapIt8nPK50j9-QU54hugRiaBh-sw,58088
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
-x_transformers-1.21.3.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.21.3.dist-info/METADATA,sha256=MX9qPfk5TEOp1xIuVqNkb3fjw-naoR_LEFt6pqQ8J88,661
-x_transformers-1.21.3.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-x_transformers-1.21.3.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.21.3.dist-info/RECORD,,
+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,,