x-transformers 2.11.4__py3-none-any.whl → 2.11.7__py3-none-any.whl

x_transformers/free_transformer.py

@@ -66,7 +66,6 @@ class BinaryMapper(Module):
 
         self.bits = bits
         self.num_codes = 2 ** bits
-        self.kl_loss_threshold = kl_loss_threshold
 
         power_two = 2 ** arange(bits)
         codes = (arange(self.num_codes)[:, None].bitwise_and(power_two) != 0).byte().bool()
@@ -74,13 +73,20 @@ class BinaryMapper(Module):
         self.register_buffer('power_two', power_two, persistent = False)
         self.register_buffer('codes', codes, persistent = False)
 
+        # aux loss
+
+        self.kl_loss_threshold = kl_loss_threshold
+        self.register_buffer('zero', tensor(0.), persistent = False)
+
     def forward(
         self,
         logits,
         temperature = 1.,
-        straight_through = None
+        straight_through = None,
+        calc_aux_loss = None
     ):
         straight_through = default(straight_through, self.training)
+        calc_aux_loss = default(calc_aux_loss, self.training)
 
         assert logits.shape[-1] == self.bits, f'logits must have a last dimension of {self.bits}'
 
@@ -95,26 +101,29 @@ class BinaryMapper(Module):
 
         one_hot = F.one_hot(indices, self.num_codes).float()
 
-        # return hard one hot if not training or overridden
+        # maybe calculate aux loss
+
+        aux_kl_loss = self.zero
 
-        if not straight_through:
-            return one_hot
+        if calc_aux_loss:
+            # calculate negative entropy
 
-        # calculate negative entropy
+            kl_div = self.bits * NAT - binary_entropy(logits)
+            aux_kl_loss = F.relu(kl_div - self.kl_loss_threshold).mean()
 
-        kl_div = self.bits * NAT - binary_entropy(logits)
-        aux_kl_loss = F.relu(kl_div - self.kl_loss_threshold).mean()
+        # maybe straight through
 
-        # get the soft G for the gradients and do a straight through
+        if straight_through:
+            # get the soft G for the gradients and do a straight through
 
-        soft_G = (
-            einsum(F.logsigmoid(logits), self.codes.float(), '... bits, codes bits -> ... codes') +
-            einsum(F.logsigmoid(-logits), (~self.codes).float(), '... bits, codes bits -> ... codes')
-        ).exp()
+            soft_G = (
+                einsum(F.logsigmoid(logits), self.codes.float(), '... bits, codes bits -> ... codes') +
+                einsum(F.logsigmoid(-logits), (~self.codes).float(), '... bits, codes bits -> ... codes')
+            ).exp()
 
-        # straight through
+            # straight through
 
-        one_hot = one_hot + soft_G - soft_G.detach()
+            one_hot = one_hot + soft_G - soft_G.detach()
 
         return one_hot, aux_kl_loss
 
@@ -161,6 +170,9 @@ class FreeTransformer(Module):
             heads = heads,
             only_cross = True,
             cross_attend = True,
+            use_rmsnorm = True,
+            rotary_pos_emb = True,
+            pre_norm_has_final_norm = True,
             **kwargs,
             **enc_kwargs
         )
@@ -180,6 +192,8 @@ class FreeTransformer(Module):
             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
@@ -190,6 +204,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
@@ -236,7 +252,7 @@ class FreeTransformer(Module):
 
         bit_logits = self.to_latent_bit_logits(pooled)
 
-        one_hot_latents, kl_loss = self.binary_mapper(bit_logits, straight_through = True)
+        one_hot_latents, kl_loss = self.binary_mapper(bit_logits, calc_aux_loss = return_kl_loss)
 
         if not return_kl_loss:
             return one_hot_latents

{x_transformers-2.11.4.dist-info → x_transformers-2.11.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.11.4
+Version: 2.11.7
 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.4.dist-info → x_transformers-2.11.7.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=t-PyeiqNq-B8BZUGICECyWCeU5XLPSGbFsQ3ICxTtsM,9072
+x_transformers/free_transformer.py,sha256=RVEMvmiSVrTJeLbXVoI4YeR26kP2mlqmwFHHLw7IJ_Q,9582
 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.4.dist-info/METADATA,sha256=18u2dW5d3aJeH3LaVsRRA_vf1gZzbSjS0Fbd5XM4dJc,96011
-x_transformers-2.11.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.11.4.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.11.4.dist-info/RECORD,,
+x_transformers-2.11.7.dist-info/METADATA,sha256=1rStN6vwvwU-0hm17zWx8yXL99ZrZ-_KGMsFwvRIPVg,96011
+x_transformers-2.11.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.11.7.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.11.7.dist-info/RECORD,,