x-transformers 2.1.35__py3-none-any.whl → 2.1.37__py3-none-any.whl

x_transformers/x_transformers.py

@@ -864,14 +864,15 @@ class DynamicTanh(Module):
         self.gamma = nn.Parameter(torch.ones(dim))
         self.beta = nn.Parameter(torch.zeros(dim))
 
-        self.unit_offset = int(unit_offset)
+        self.pre_tanh_scale_offset = init_alpha if unit_offset else 0.
+        self.gamma_offset = float(unit_offset)
 
-        nn.init.constant_(self.pre_tanh_scale, 1. - float(unit_offset))
+        nn.init.constant_(self.pre_tanh_scale, 0 if unit_offset else init_alpha)
         nn.init.constant_(self.gamma, 1. - float(unit_offset))
 
     def forward(self, x):
-        pre_tanh_scale = self.pre_tanh_scale + self.unit_offset
-        gamma = self.gamma + self.unit_offset
+        pre_tanh_scale = self.pre_tanh_scale + self.pre_tanh_scale_offset
+        gamma = self.gamma + self.gamma_offset
         return (x * pre_tanh_scale).tanh() * gamma + self.beta
 
 # residual and residual gates
@@ -2908,6 +2909,7 @@ class TransformerWrapper(Module):
         return_embeddings = False,
         return_logits_and_embeddings = False,
         return_intermediates = False,
+        return_embeddings_and_intermediates = False,
         return_logit_entropies = False,
         mask = None,
         return_mems = False,
@@ -2939,8 +2941,8 @@ class TransformerWrapper(Module):
 
         b, n, device, num_mems, has_memory_tokens, emb_frac_gradient, orig_mask = x.shape[0], x.shape[1], x.device, self.num_memory_tokens, self.num_memory_tokens > 0, self.emb_frac_gradient, mask
 
-        return_hiddens = return_mems | return_attn | return_intermediates | return_attn_z_loss
-        return_embeddings = return_embeddings | (not exists(self.to_logits))
+        return_hiddens = return_mems | return_attn | return_intermediates | return_attn_z_loss | return_embeddings_and_intermediates
+        return_embeddings = return_embeddings | (not exists(self.to_logits)) | return_embeddings_and_intermediates
 
         # absolute positional embedding
 
@@ -3130,6 +3132,8 @@ class TransformerWrapper(Module):
 
         if return_logits_and_embeddings:
             out = (logits, x)
+        elif return_embeddings_and_intermediates:
+            out = (x, intermediates)
         elif return_embeddings:
             out = x
         else:

{x_transformers-2.1.35.dist-info → x_transformers-2.1.37.dist-info}/METADATA

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

{x_transformers-2.1.35.dist-info → x_transformers-2.1.37.dist-info}/RECORD

@@ -7,10 +7,10 @@ x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
-x_transformers/x_transformers.py,sha256=oyYk31qhDyt6cCuDeaHBl4XhUo5bfIwmYE_z1E1OpXU,111502
+x_transformers/x_transformers.py,sha256=twoqq2kfVWxntitHKLs2sxFMhK1CPLxGGBDAmkiHXcM,111812
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-2.1.35.dist-info/METADATA,sha256=tLbl-c1QtaOphTa1DpdNfh4dXzFwTt9Fvdh94tnwdTs,88161
-x_transformers-2.1.35.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.1.35.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.1.35.dist-info/RECORD,,
+x_transformers-2.1.37.dist-info/METADATA,sha256=uaCIy-GAGH4OPrYa0mxjJJ-FDtMlMuiIbg1sQPb3BRw,88161
+x_transformers-2.1.37.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.1.37.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.1.37.dist-info/RECORD,,