x-transformers 1.23.2__py3-none-any.whl → 1.23.3__py3-none-any.whl

x_transformers/x_transformers.py

@@ -1622,6 +1622,7 @@ class ContinuousTransformerWrapper(nn.Module):
         dim_out = None,
         emb_dim = None,
         max_mem_len = 0,
+        num_memory_tokens = None,
         post_emb_norm = False,
         emb_dropout = 0.,
         use_abs_pos_emb = True,
@@ -1646,10 +1647,21 @@ class ContinuousTransformerWrapper(nn.Module):
         self.post_emb_norm = nn.LayerNorm(dim) if post_emb_norm else nn.Identity()
         self.emb_dropout = nn.Dropout(emb_dropout)
 
-        self.project_in = nn.Linear(dim_in, dim) if exists(dim_in) else nn.Identity()
+        # memory tokens
+
+        num_memory_tokens = default(num_memory_tokens, 0)
+        self.has_memory_tokens = num_memory_tokens > 0
+
+        if num_memory_tokens > 0:
+            self.memory_tokens = nn.Parameter(torch.randn(num_memory_tokens, dim))
+
+        # attention layers
 
         self.attn_layers = attn_layers
 
+        # project in and out
+
+        self.project_in = nn.Linear(dim_in, dim) if exists(dim_in) else nn.Identity()
         self.project_out = nn.Linear(dim, dim_out) if exists(dim_out) else nn.Identity()
 
     def forward(
@@ -1665,11 +1677,19 @@ class ContinuousTransformerWrapper(nn.Module):
         prepend_embeds = None,
         **kwargs
     ):
+        batch = x.shape[0]
+
         x = self.project_in(x)
         x = x + self.pos_emb(x, pos = pos)
 
         x = self.post_emb_norm(x)
 
+        # memory tokens
+
+        if self.has_memory_tokens:
+            m = repeat(self.memory_tokens, 'm d -> b m d', b = batch)
+            x, mem_ps = pack([m, x], 'b * d')
+
         # whether to append embeds, as in PaLI, for image embeddings
 
         if exists(prepend_embeds):
@@ -1680,8 +1700,15 @@ class ContinuousTransformerWrapper(nn.Module):
 
         x = self.emb_dropout(x)
 
+        # attention layers
+
         x, intermediates = self.attn_layers(x, mask = mask, mems = mems, return_hiddens = True, **kwargs)
 
+        # splice out memory tokens
+
+        if self.has_memory_tokens:
+            m, x = unpack(x, mem_ps, 'b * d')
+
         out = self.project_out(x) if not return_embeddings else x
 
         if return_intermediates:

{x_transformers-1.23.2.dist-info → x_transformers-1.23.3.dist-info}/METADATA

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

{x_transformers-1.23.2.dist-info → x_transformers-1.23.3.dist-info}/RECORD

@@ -3,10 +3,10 @@ x_transformers/attend.py,sha256=hZcz_iijzbEqbXp2_BPEVL-1LoHXmYaHE6e6Oy-7hFE,1126
 x_transformers/autoregressive_wrapper.py,sha256=f2u0usjUfAlXwgTz87O8J8XjGTbsbrx2XEP6K2beSNI,8944
 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=KQ9mU_jE27whl6yQI67grF0S8Xhd3GndnM6Yd0-q-lw,61162
+x_transformers/x_transformers.py,sha256=goudsIa79mfyJtzuI0GqTSdGQ5CXG1ga5Is9h3UBC5Y,61861
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
-x_transformers-1.23.2.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.23.2.dist-info/METADATA,sha256=8h0sbx8-4yNTOJuAZLbe5HQ16hsmZI1M_mT-rMIIMJc,661
-x_transformers-1.23.2.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-x_transformers-1.23.2.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.23.2.dist-info/RECORD,,
+x_transformers-1.23.3.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.23.3.dist-info/METADATA,sha256=SXNDjqYSGkklnbXVRg8S52VxDR6VVO62KvRH60abY_k,661
+x_transformers-1.23.3.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
+x_transformers-1.23.3.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.23.3.dist-info/RECORD,,