x-transformers 1.42.3__py3-none-any.whl → 1.42.5__py3-none-any.whl

x_transformers/x_transformers.py

@@ -512,12 +512,15 @@ class DataDependentAlibi(Module):
         self,
         dim,
         heads,
+        causal = True,
         bias_init = 5.,
-        post_log_scale = 1.
+        post_log_scale = 1.,
     ):
         super().__init__()
 
-        linear = nn.Linear(dim, heads)
+        self.causal = causal
+
+        linear = nn.Linear(dim, heads * (1 if causal else 2))
 
         self.to_forget_gates = nn.Sequential(
             linear,
@@ -529,9 +532,21 @@ class DataDependentAlibi(Module):
         self.post_log_scale = post_log_scale
 
     def forward(self, x):
+        bidirectional = not self.causal
+
         forget_gates = self.to_forget_gates(x) * self.post_log_scale
+
         forget_gates = forget_gates.cumsum(dim = -1)
+
+        if bidirectional:
+            forget_gates, forget_gates_reversed = forget_gates.chunk(2, dim = 1)
+
         forget_gates = einx.subtract('b h i, b h j -> b h i j', forget_gates, forget_gates)
+
+        if bidirectional:
+            forget_gates_reversed = einx.subtract('b h j, b h i -> b h i j', forget_gates_reversed, forget_gates_reversed)
+            forget_gates = forget_gates.tril() + forget_gates_reversed.triu()
+
         return forget_gates
 
 class PerRowDataDependentAlibi(Module):
@@ -541,10 +556,13 @@ class PerRowDataDependentAlibi(Module):
         self,
         dim,
         heads,
+        causal = True,
         dim_head = 8,
         post_log_scale = 1.
     ):
         super().__init__()
+        assert causal, 'bidirectional not supported yet'
+
         self.scale = dim_head ** -0.5
 
         linear = nn.Linear(dim, heads * dim_head * 2, bias = False)
@@ -1138,10 +1156,9 @@ class Attention(Module):
         self.data_dependent_alibi = None
 
         if data_dependent_alibi:
-            assert causal, 'data dependent alibi only works for autoregressive for now until further research'
 
             dda_klass = DataDependentAlibi if not data_dependent_alibi_per_row else PerRowDataDependentAlibi
-            dda_kwargs = dict(dim = dim, heads = heads)
+            dda_kwargs = dict(dim = dim, heads = heads, causal = causal)
 
             if data_dependent_alibi_per_row:
                 dda_kwargs.update(dim_head = data_dependent_alibi_per_row_dim_head)
@@ -1252,6 +1269,20 @@ class Attention(Module):
 
         k, v, r = tuple(maybe(rearrange)(t, 'b n (h d) -> b h n d', h = kv_h) for t in (k, v, r))
 
+        # if previous values passed in for residual, either invoke resformer or neutreno
+
+        orig_values = v
+
+        if exists(value_residual):
+            if self.neutreno_value_residual:
+                diff_values = (value_residual - v) * self.neutreno_alpha
+                diff_values = repeat(diff_values, 'b h n d -> b (r h) n d', r = h // kv_h)
+            else:
+                # https://arxiv.org/abs/2410.17897v1
+                v = 0.5 * (v + value_residual)
+
+        # take care of caching
+
         if exists(cache):
             ck, cv = cache.cached_kv
 
@@ -1363,16 +1394,6 @@ class Attention(Module):
         if exists(self.data_dependent_alibi):
             attn_bias = self.data_dependent_alibi(x)
 
-        # if previous values passed in for residual, either invoke resformer or neutreno
-
-        if exists(value_residual):
-            if self.neutreno_value_residual:
-                diff_values = (value_residual - v) * self.neutreno_alpha
-                diff_values = repeat(diff_values, 'b h n d -> b (r h) n d', r = h // kv_h)
-            else:
-                # https://arxiv.org/abs/2410.17897v1
-                v = 0.5 * (v + value_residual)
-
         # attention is all we need
 
         out, intermediates = self.attend(
@@ -1384,7 +1405,7 @@ class Attention(Module):
 
         # store the values for resformer or Neutreno
 
-        intermediates.values = v
+        intermediates.values = orig_values
 
         if exists(value_residual) and self.neutreno_value_residual:
             out = out + diff_values

{x_transformers-1.42.3.dist-info → x_transformers-1.42.5.dist-info}/METADATA

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

{x_transformers-1.42.3.dist-info → x_transformers-1.42.5.dist-info}/RECORD

@@ -6,11 +6,11 @@ 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=K8zt6n7aC8iMMrWJ-0ryNsJeq7eXR94ci5DXAElC-lY,91995
+x_transformers/x_transformers.py,sha256=KIR7efx59xl0BVshU1e6RO0YKgz7zeYBXITDNYWJ4mQ,92506
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-1.42.3.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.42.3.dist-info/METADATA,sha256=A9f9ByY6e62AuBpOEMBndE_b_8TOovaau75D2fo73H4,689
-x_transformers-1.42.3.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-x_transformers-1.42.3.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.42.3.dist-info/RECORD,,
+x_transformers-1.42.5.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.42.5.dist-info/METADATA,sha256=cLnay5nt6F6GKdghsaqHiZQsVmJ9dS5l-IDozZIs3ec,689
+x_transformers-1.42.5.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+x_transformers-1.42.5.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.42.5.dist-info/RECORD,,