x-transformers 1.32.7__py3-none-any.whl → 1.32.8__py3-none-any.whl

x_transformers/continuous.py

@@ -9,6 +9,7 @@ from x_transformers.x_transformers import (
     ScaledSinusoidalEmbedding,
     AbsolutePositionalEmbedding,
     LayerNorm,
+    masked_mean,
     always,
     pad_at_dim
 )
@@ -39,7 +40,8 @@ class ContinuousTransformerWrapper(nn.Module):
         post_emb_norm = False,
         emb_dropout = 0.,
         use_abs_pos_emb = True,
-        scaled_sinu_pos_emb = False
+        scaled_sinu_pos_emb = False,
+        average_pool_embed = False
     ):
         super().__init__()
         dim = attn_layers.dim
@@ -72,6 +74,10 @@ class ContinuousTransformerWrapper(nn.Module):
 
         self.attn_layers = attn_layers
 
+        # average pool
+
+        self.average_pool_embed = average_pool_embed
+
         # project in and out
 
         self.project_in = nn.Linear(dim_in, dim, bias = False) if exists(dim_in) else nn.Identity()
@@ -92,7 +98,7 @@ class ContinuousTransformerWrapper(nn.Module):
         prepend_mask = None,
         **kwargs
     ):
-        batch, seq, device = *x.shape[:2], x.device
+        batch, seq, orig_mask, device = *x.shape[:2], mask, x.device
 
         x = self.project_in(x)
         x = x + self.pos_emb(x, pos = pos)
@@ -136,6 +142,11 @@ class ContinuousTransformerWrapper(nn.Module):
             m, x = unpack(x, mem_ps, 'b * d')
             intermediates.memory_tokens = m
 
+        if self.average_pool_embed:
+            x = masked_mean(x, mask = orig_mask)
+
+        # maybe linear project out
+
         out = self.project_out(x) if not return_embeddings else x
 
         if return_intermediates:

x_transformers/x_transformers.py

@@ -99,6 +99,17 @@ def l2norm(t, groups = 1):
 def softclamp(t, value):
     return (t / value).tanh() * value
 
+def masked_mean(t, mask = None, dim = 1):
+    if not exists(mask):
+        return t.mean(dim = dim)
+
+    dims_append = (1,) * (t.ndim - mask.ndim)
+    mask = mask.reshape(*mask.shape, *dims_append)
+
+    num = (t * mask).sum(dim = dim)
+    den = mask.sum(dim = dim).clamp(min = 1.)
+    return num / den
+
 def pad_at_dim(t, pad: Tuple[int, int], dim = -1, value = 0.):
     if pad == (0, 0):
         return t
@@ -1666,6 +1677,8 @@ class AttentionLayers(Module):
 
         layer_variables = tuple(tuple(layer_variable[i] for i in layers_execute_order) for layer_variable in layer_variables)
 
+        first_skip = None
+
         # go through the attention and feedforward layers
 
         for ind, (layer_type, (norm, block, residual_fn), layer_dropout) in enumerate(zip(*layer_variables)):
@@ -1674,6 +1687,9 @@ class AttentionLayers(Module):
             if self.training and layer_dropout > 0. and random() < layer_dropout:
                 continue
 
+            if ind == 1:
+                first_skip = x.clone()
+
             if layer_type == 'a':
                 if return_hiddens:
                     hiddens.append(x)
@@ -1909,6 +1925,7 @@ class TransformerWrapper(Module):
         l2norm_embed = False,
         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
     ):
         super().__init__()
 
@@ -1954,6 +1971,10 @@ class TransformerWrapper(Module):
 
         assert num_output_heads > 0
 
+        # whether to average pool the embed (`global average pool`)
+
+        self.average_pool_embed = average_pool_embed
+
         # output head, usually to logits of num_tokens
 
         logits_dim = default(logits_dim, num_tokens)
@@ -2015,7 +2036,7 @@ class TransformerWrapper(Module):
         cache: LayerIntermediates | None = None,
         **kwargs
     ):
-        b, n, device, num_mems, has_memory_tokens, emb_frac_gradient = x.shape[0], x.shape[1], x.device, self.num_memory_tokens, self.num_memory_tokens > 0, self.emb_frac_gradient
+        b, n, device, num_mems, has_memory_tokens, emb_frac_gradient, orig_mask = x.shape[0], x.shape[1], x.device, self.num_memory_tokens, self.num_memory_tokens > 0, self.emb_frac_gradient, mask
 
         return_hiddens = return_mems | return_attn | return_intermediates | return_attn_z_loss
         return_embeddings = return_embeddings | (not exists(self.to_logits))
@@ -2118,6 +2139,11 @@ class TransformerWrapper(Module):
 
             x = x[:, :n]
 
+        # global average pool
+
+        if self.average_pool_embed:
+            x = masked_mean(x, mask = orig_mask, dim = 1)
+
         # projecting to logits
 
         if not return_embeddings:

{x_transformers-1.32.7.dist-info → x_transformers-1.32.8.dist-info}/METADATA

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

{x_transformers-1.32.7.dist-info → x_transformers-1.32.8.dist-info}/RECORD

@@ -1,15 +1,15 @@
 x_transformers/__init__.py,sha256=-MkQrSc37cTVDX7AOykxunYnqVtFlQ7lb0Cse5dsGWU,793
 x_transformers/attend.py,sha256=MI-m91wumBFqFqr_KK9MLgsLk_vPeaVbFMyDr_mWdmY,11349
 x_transformers/autoregressive_wrapper.py,sha256=uX8Mb0zLsQrZECt_9UGt35g7tC05Rk3nPqO6xp2FFCc,9619
-x_transformers/continuous.py,sha256=WO52n9lFAXv5-SGadi2cApGF8dkouN8QSTEOuC7erj8,6180
+x_transformers/continuous.py,sha256=cIVEdhfei258__ziV7kQBrJMxCel54bExBTDrO9rfCI,6450
 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=5DHbYgx0RPg9QHvfBs2qHWrtn4Jji-q0d1MRBbcRPR8,76696
+x_transformers/x_transformers.py,sha256=cq364zjUVvGEeFxdu703yI2tp1VhpxTIpLTgMshHpzI,77392
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-x_transformers-1.32.7.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.32.7.dist-info/METADATA,sha256=25J9CJ3OxsR_SZkvubPhyjSN-NmvU_yVVQHNMFzoKVg,661
-x_transformers-1.32.7.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-x_transformers-1.32.7.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.32.7.dist-info/RECORD,,
+x_transformers-1.32.8.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.32.8.dist-info/METADATA,sha256=tEHQVjqqXKQ2eSd-j5yrmfubkjuQZga_sT_5XgHanQo,661
+x_transformers-1.32.8.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+x_transformers-1.32.8.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.32.8.dist-info/RECORD,,