x-transformers 1.41.2__py3-none-any.whl → 1.41.3__py3-none-any.whl

x_transformers/x_transformers.py

@@ -499,7 +499,9 @@ class DataDependentAlibi(Module):
     def __init__(
         self,
         dim,
-        heads
+        heads,
+        bias_init = 5.,
+        post_log_scale = 1.
     ):
         super().__init__()
 
@@ -511,22 +513,24 @@ class DataDependentAlibi(Module):
             nn.LogSigmoid()
         )
 
-        nn.init.constant_(linear.bias, 5.)
+        nn.init.constant_(linear.bias, bias_init)
+        self.post_log_scale = post_log_scale
 
     def forward(self, x):
-        forget_gates = self.to_forget_gates(x)
+        forget_gates = self.to_forget_gates(x) * self.post_log_scale
         forget_gates = forget_gates.cumsum(dim = -1)
         forget_gates = einx.subtract('b h i, b h j -> b h i j', forget_gates, forget_gates)
         return forget_gates
 
 class PerRowDataDependentAlibi(Module):
-    """ same as data dependent alibi from forgetting transformer, but the forgetting gates are also derived by a queris and keys with a small head dimension """
+    """ same as data dependent alibi from forgetting transformer, but the forgetting gates are also derived by a queries and keys with a small head dimension """
 
     def __init__(
         self,
         dim,
         heads,
-        dim_head = 8
+        dim_head = 8,
+        post_log_scale = 1.
     ):
         super().__init__()
         self.scale = dim_head ** -0.5
@@ -535,14 +539,16 @@ class PerRowDataDependentAlibi(Module):
 
         self.to_forget_gates = nn.Sequential(
             linear,
-            Rearrange('b n (kv h d) -> kv b h n d', kv = 2, d = dim_head)
+            Rearrange('b n (qk h d) -> qk b h n d', qk = 2, d = dim_head)
         )
 
+        self.post_log_scale = post_log_scale
+
     def forward(self, x):
         q, k = self.to_forget_gates(x)
         forget_gates = einsum('... i d, ... j d -> ... i j', q, k) * self.scale
 
-        forget_gates = F.logsigmoid(forget_gates)
+        forget_gates = F.logsigmoid(forget_gates) * self.post_log_scale
 
         # mask out upper triangle + diagonal
 
@@ -1010,6 +1016,7 @@ class Attention(Module):
         data_dependent_alibi = False,
         data_dependent_alibi_per_row = False,
         data_dependent_alibi_per_row_dim_head = 8,
+        data_dependent_alibi_kwargs: dict = dict(),
         use_cope = False,
         cope_max_pos = 16,
         cope_soft_onehot_pos = False,
@@ -1127,7 +1134,7 @@ class Attention(Module):
             if data_dependent_alibi_per_row:
                 dda_kwargs.update(dim_head = data_dependent_alibi_per_row_dim_head)
 
-            self.data_dependent_alibi = dda_klass(**dda_kwargs)
+            self.data_dependent_alibi = dda_klass(**dda_kwargs, **data_dependent_alibi_kwargs)
 
         # attend class - includes core attention algorithm + talking heads
 

{x_transformers-1.41.2.dist-info → x_transformers-1.41.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.41.2
+Version: 1.41.3
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang

{x_transformers-1.41.2.dist-info → x_transformers-1.41.3.dist-info}/RECORD

@@ -5,11 +5,11 @@ x_transformers/continuous.py,sha256=cIVEdhfei258__ziV7kQBrJMxCel54bExBTDrO9rfCI,
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
-x_transformers/x_transformers.py,sha256=5yoN96mTpw3qygZg2plVed5lTOs9p0REtt50lOlllHU,91334
+x_transformers/x_transformers.py,sha256=aaxMw4iJkLAh8D4W0g8EwWDfIJVPYdhpmj5T9ZFa0js,91642
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-1.41.2.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.41.2.dist-info/METADATA,sha256=84WHjWdVHIxkHdHxMEKK5OEssLU80cdj4aGTJQ7SpO4,689
-x_transformers-1.41.2.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-x_transformers-1.41.2.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.41.2.dist-info/RECORD,,
+x_transformers-1.41.3.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.41.3.dist-info/METADATA,sha256=GyMdKP9ErEnGHbBqY3MwpSGSE41lK_KQbeE5n_O2Py4,689
+x_transformers-1.41.3.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+x_transformers-1.41.3.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.41.3.dist-info/RECORD,,