x-transformers 2.11.15__py3-none-any.whl → 2.11.16__py3-none-any.whl

x_transformers/free_transformer.py

@@ -225,8 +225,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 +238,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 +345,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 +360,25 @@ 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)
 
         condition = self.from_latent_to_condition(latents)
 

{x_transformers-2.11.15.dist-info → x_transformers-2.11.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.11.15
+Version: 2.11.16
 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.dist-info → x_transformers-2.11.16.dist-info}/RECORD

@@ -5,7 +5,7 @@ x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTN
 x_transformers/continuous.py,sha256=WwpQCjyVY4PtuEAOFY68zqgklbF9I7AL5w6874YlDe8,13249
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/entropy_based_tokenizer.py,sha256=F2lO8-v3aLIcVDVNhu7RR-UtRdlmaaYJzBK9m7OnLE8,5018
-x_transformers/free_transformer.py,sha256=a_sF_tx2RgKNsPCum22jlYam28OWEk3B0o1D4-Vo9Fw,10714
+x_transformers/free_transformer.py,sha256=_hYYkaro3xei3MC3rwtuCWi9gSnciXyAT91_7SrA0nw,11396
 x_transformers/gpt_vae.py,sha256=1zyjwgfZr6CRDsh5VMCPSdoCPg-sdX5mXmZ_mn4VyYQ,6082
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
@@ -14,7 +14,7 @@ x_transformers/up_wrapper.py,sha256=YC2LN14_7Xx9Wtiek2rtEJ_qHqdfSmKlh3d7Cgxwd80,
 x_transformers/x_transformers.py,sha256=5ctPu8tvlbUMrtW360e_LPnoGv6xcgQFsyWdbvLo6Tk,127002
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.11.15.dist-info/METADATA,sha256=DSY5ug0mmywhOVxsCxjVkIzyWNY9ot4kmUxBFresdaE,96012
-x_transformers-2.11.15.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.11.15.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.11.15.dist-info/RECORD,,
+x_transformers-2.11.16.dist-info/METADATA,sha256=cvhm5LnIRCdqLuv25iSU4vj0a6Np9j2lv2O9W-V48-k,96012
+x_transformers-2.11.16.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.11.16.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.11.16.dist-info/RECORD,,