x-transformers 2.11.15__tar.gz → 2.11.17__tar.gz

{x_transformers-2.11.15 → x_transformers-2.11.17}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.11.15
+Version: 2.11.17
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers

{x_transformers-2.11.15 → x_transformers-2.11.17}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-transformers"
-version = "2.11.15"
+version = "2.11.17"
 description = "X-Transformers"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{x_transformers-2.11.15 → x_transformers-2.11.17}/tests/test_x_transformers.py

@@ -1411,9 +1411,11 @@ def test_attn_negative_weights(
 
 @param('per_token_latents', (False, True))
 @param('dec_head_depth', (0, 4))
+@param('separate_seq_for_latents', (False, True))
 def test_free(
     dec_head_depth,
-    per_token_latents
+    per_token_latents,
+    separate_seq_for_latents
 ):
     from x_transformers.free_transformer import FreeTransformer
 
@@ -1432,7 +1434,9 @@ def test_free(
 
     seq = torch.randint(0, 256, (1, 1024))
 
-    loss, (ar_loss, aux_loss) = model(seq, return_all_losses = True)
+    separate_seq_for_latents = torch.randint(0, 256, (1, 32)) if separate_seq_for_latents else None
+
+    loss, (ar_loss, aux_loss) = model(seq, separate_seq_for_latents, return_all_losses = True)
     loss.backward()
 
     assert aux_loss.numel() == 1

{x_transformers-2.11.15 → x_transformers-2.11.17}/x_transformers/free_transformer.py

@@ -149,6 +149,7 @@ class FreeTransformer(Module):
         enc_kwargs: dict = dict(),
         dec_kwargs: dict = dict(),
         kl_loss_weight = 1.,
+        latent_dropout_prob = 0.,
         pad_id = -1,
         **kwargs
     ):
@@ -187,6 +188,8 @@ class FreeTransformer(Module):
 
         self.from_latent_to_condition = nn.Linear(self.binary_mapper.num_codes, dim, bias = False)
 
+        self.latent_dropout = nn.Dropout(latent_dropout_prob)
+
         self.decoder_head = Decoder(
             dim = dim,
             depth = dec_head_depth,
@@ -225,8 +228,11 @@ class FreeTransformer(Module):
         self,
         decoder_head_embeds,
         mask = None,
-        return_kl_loss = False
+        return_kl_loss = False,
+        per_token_latents = None
     ):
+        per_token_latents = default(per_token_latents, self.per_token_latents)
+
         batch, seq_len, device = *decoder_head_embeds.shape[:2], decoder_head_embeds.device
 
         query_tokens = repeat(self.query_token_for_latents, 'd -> b 1 d', b = batch)
@@ -235,7 +241,7 @@ class FreeTransformer(Module):
 
         # handle the interesting per query token latents, as in the paper
 
-        if self.per_token_latents:
+        if per_token_latents:
             query_tokens = repeat(query_tokens, 'b 1 d -> b n d', n = seq_len)
 
             rotary_pos = torch.arange(seq_len, device = device)
@@ -342,13 +348,13 @@ class FreeTransformer(Module):
     def forward(
         self,
         seq,
+        seq_for_latents = None,
         return_all_losses = False
     ):
         batch, device = seq.shape[0], seq.device
 
         seq, labels = seq[:, :-1], seq[:, 1:]
 
-        encoder_mask = seq != self.pad_id
 
         tokens = self.token_emb(seq)
 
@@ -357,9 +363,27 @@ class FreeTransformer(Module):
         if exists(self.decoder_head):
             tokens = self.decoder_head(tokens)
 
+        # determine whether to use a separate sequence for encoding latents
+
+        if exists(seq_for_latents):
+            tokens_for_latents = self.token_emb(seq_for_latents)
+
+            if exists(self.decoder_head):
+                tokens_for_latents = self.decoder_head(tokens_for_latents)
+
+            encoder_mask = seq_for_latents != self.pad_id
+            per_token_latents = False
+        else:
+
+            tokens_for_latents = tokens
+            encoder_mask = seq != self.pad_id
+            per_token_latents = None
+
         # get latent Z
 
-        latents, kl_loss = self.encode_to_latents(tokens, mask = encoder_mask, return_kl_loss = True)
+        latents, kl_loss = self.encode_to_latents(tokens_for_latents, mask = encoder_mask, per_token_latents = per_token_latents, return_kl_loss = True)
+
+        latents = self.latent_dropout(latents)
 
         condition = self.from_latent_to_condition(latents)