x-transformers 2.11.2__py3-none-any.whl → 2.11.5__py3-none-any.whl

x_transformers/free_transformer.py

@@ -128,19 +128,19 @@ class FreeTransformer(Module):
         dim,
         dec_head_depth,
         dec_tail_depth,
-        enc_depth,
         max_seq_len,
+        enc_depth = 1,
         dim_latent = None,
         attn_dim_head = 64,
         heads = 8,
         latent_bits = 16,
+        per_token_latents = True,  # they use a latent per token in the sequence, instead of one for entire sequence, iiuc
         kl_loss_threshold = NAT,
         binary_mapper_kwargs: dict = dict(),
         enc_kwargs: dict = dict(),
         dec_kwargs: dict = dict(),
         kl_loss_weight = 1.,
         pad_id = -1,
-        encoder: Module | None = None,
         **kwargs
     ):
         super().__init__()
@@ -150,39 +150,40 @@ class FreeTransformer(Module):
 
         self.token_unembed = nn.Linear(dim, num_tokens, bias = False)
 
-        if not exists(encoder):
-            encoder = Encoder(
-                dim = dim,
-                depth = enc_depth,
-                attn_dim_head = attn_dim_head,
-                heads = heads,
-                **kwargs,
-                **enc_kwargs
-            )
+        self.query_token_for_latents = nn.Parameter(torch.randn(dim) * 1e-2)
 
-        self.encoder = encoder
+        self.per_token_latents = per_token_latents
 
-        self.to_latent_bit_logits = nn.Sequential(
-            Reduce('b n d -> b d', 'mean'),
-            nn.Linear(dim, latent_bits, bias = False),
+        self.encoder = Encoder(
+            dim = dim,
+            depth = enc_depth,
+            attn_dim_head = attn_dim_head,
+            heads = heads,
+            only_cross = True,
+            cross_attend = True,
+            use_rmsnorm = True,
+            rotary_pos_emb = True,
+            **kwargs,
+            **enc_kwargs
         )
 
+        self.to_latent_bit_logits = nn.Linear(dim, latent_bits, bias = False)
+
         self.binary_mapper = BinaryMapper(
             latent_bits,
             kl_loss_threshold,
             **binary_mapper_kwargs
         )
 
-        self.from_latent_to_condition = nn.Sequential(
-            nn.Linear(2 ** latent_bits, dim, bias = False),
-            Rearrange('b d -> b 1 d')
-        )
+        self.from_latent_to_condition = nn.Linear(self.binary_mapper.num_codes, dim, bias = False)
 
         self.decoder_head = Decoder(
             dim = dim,
             depth = dec_head_depth,
             attn_dim_head = attn_dim_head,
             heads = heads,
+            rotary_pos_emb = True,
+            use_rmsnorm = True,
             pre_norm_has_final_norm = False,
             **kwargs,
             **dec_kwargs
@@ -193,6 +194,8 @@ class FreeTransformer(Module):
             depth = dec_tail_depth,
             attn_dim_head = attn_dim_head,
             heads = heads,
+            rotary_pos_emb = True,
+            use_rmsnorm = True,
             pre_norm_has_final_norm = True,
             **kwargs,
             **dec_kwargs
@@ -208,11 +211,34 @@ class FreeTransformer(Module):
 
     def encode_to_latents(
         self,
-        seq,
+        decoder_head_embeds,
         mask = None,
         return_kl_loss = False
     ):
-        pooled = self.encoder(seq, mask = mask)
+        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)
+
+        encoder_kwargs = dict()
+
+        # handle the interesting per query token latents, as in the paper
+
+        if self.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)
+
+            encoder_kwargs.update(
+                pos = rotary_pos,
+                context_pos = rotary_pos
+            )
+
+        pooled = self.encoder(
+            query_tokens,
+            context = decoder_head_embeds,
+            context_mask = mask,
+            **encoder_kwargs
+        )
 
         bit_logits = self.to_latent_bit_logits(pooled)
 

{x_transformers-2.11.2.dist-info → x_transformers-2.11.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.11.2
+Version: 2.11.5
 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.2.dist-info → x_transformers-2.11.5.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=K9dX9xj0pJgiOi9jlOLCD9Nn9eYNNgTWj9YvQLhexHw,8295
+x_transformers/free_transformer.py,sha256=oUrTl3cNbal9-c2wU7Q3B8vruWg_FIbFNBhsgkBryc4,9273
 x_transformers/gpt_vae.py,sha256=4QdznXZcU7pmMXUeEocAOKpcTkREYS-zDHktN5ADtNk,5981
 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=ADr83Fz2cehj_F7N1bMwxhAg-r48fGhlaZqw3hxoxMQ,125765
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.11.2.dist-info/METADATA,sha256=oDAvV32UGA_1ajVnH8qVo2lAt0tC9rMPI8ssKzA6IVE,96011
-x_transformers-2.11.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.11.2.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.11.2.dist-info/RECORD,,
+x_transformers-2.11.5.dist-info/METADATA,sha256=6lQ4z5dCmDbPAc1Vhl1ObzSqPUneWTqwYKJKaHhbiNY,96011
+x_transformers-2.11.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.11.5.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.11.5.dist-info/RECORD,,