x-transformers 2.6.5__tar.gz → 2.6.7__tar.gz

{x_transformers-2.6.5 → x_transformers-2.6.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.6.5
+Version: 2.6.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.6.5 → x_transformers-2.6.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-transformers"
-version = "2.6.5"
+version = "2.6.7"
 description = "X-Transformers"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{x_transformers-2.6.5 → x_transformers-2.6.7}/tests/test_x_transformers.py

@@ -1245,10 +1245,45 @@ def test_learned_head_attn_sink():
             dim = 512,
             depth = 12,
             heads = 8,
-            attn_head_learned_sinks = 4
+            attn_head_learned_sink = True
         )
     )
 
     seq = torch.randint(0, 20000, (3, 1024))
 
     logits = model(seq)
+
+def test_accept_layer_intermediates():
+    from x_transformers import TransformerWrapper, Decoder, AutoregressiveWrapper
+
+    vlm = TransformerWrapper(
+        num_tokens = 20000,
+        max_seq_len = 1024,
+        attn_layers = Decoder(
+            dim = 512,
+            depth = 3,
+            heads = 4,
+        )
+    )
+
+    seq = torch.randint(0, 20000, (3, 1024))
+    mask = torch.randint(0, 2, (3, 1024)).bool()
+
+    _, intermediates = vlm(seq, return_intermediates = True)
+
+    action_model = Decoder(
+        dim = 512,
+        depth = 6,
+        heads = 8,
+    )
+
+    seq = torch.randn(3, 32, 512)
+
+    embeds = action_model(
+        seq,
+        self_attn_additional_kv = intermediates,
+        detach_additional_kv = True,
+        additional_kv_mask = mask
+    )
+
+    assert embeds.shape == (3, 32, 512)

{x_transformers-2.6.5 → x_transformers-2.6.7}/x_transformers/attend.py

@@ -23,7 +23,7 @@ class Intermediates:
     pre_softmax_attn:   Tensor | None = None
     post_softmax_attn:  Tensor | None = None
     values:             Tensor | None = None
-    cached_kv:          Tuple[Tensor, Tensor] | None = None
+    cached_kv:          tuple[Tensor, Tensor] | None = None
     layer_type:         str | None = None
     hybrid_hidden:      Tensor | None = None
 
@@ -176,7 +176,7 @@ class Attend(Module):
         softclamp_logits = False,
         logit_softclamp_value = 50.,
         add_zero_kv = False,
-        head_learned_sinks = 0,
+        head_learned_sink = False,
         selective = False,
         hard = False,
         cope = None,
@@ -257,10 +257,10 @@ class Attend(Module):
 
         # learned sink concatted pre-softmax, working solution from gpt-oss
 
-        self.has_head_learned_sinks = head_learned_sinks > 0
-        assert not (self.has_head_learned_sinks and flash), f'not supported for flash attention yet'
+        assert not (head_learned_sink and flash), f'not supported for flash attention yet'
 
-        self.head_attn_sinks = Parameter(torch.zeros(heads, head_learned_sinks)) if self.has_head_learned_sinks else None
+        self.head_learned_sink = head_learned_sink
+        self.head_attn_sink = Parameter(torch.zeros(heads)) if head_learned_sink else None
 
         # soft clamp attention logit value
 
@@ -517,10 +517,9 @@ class Attend(Module):
         if self.selective:
             sim = selective_attn(sim)
 
-        if self.has_head_learned_sinks:
+        if self.head_learned_sink:
             # add learned attention sink
-            num_sinks = self.head_attn_sinks.shape[-1]
-            attn_sink = repeat(self.head_attn_sinks, 'h sinks -> b h i sinks', b = sim.shape[0], i = sim.shape[2])
+            attn_sink = repeat(self.head_attn_sink, 'h -> b h i 1', b = sim.shape[0], i = sim.shape[2])
             sim = cat((attn_sink, sim), dim = -1)
 
         pre_softmax_attn = sim
@@ -531,9 +530,9 @@ class Attend(Module):
 
         post_softmax_attn = attn
 
-        if self.has_head_learned_sinks:
+        if self.head_learned_sink:
             # remove attention sink
-            attn = attn[..., num_sinks:]
+            attn = attn[..., 1:]
 
         attn = self.attn_dropout(attn)
 

{x_transformers-2.6.5 → x_transformers-2.6.7}/x_transformers/x_transformers.py

@@ -10,7 +10,7 @@ import torch
 from torch.amp import autocast
 import torch.nn.functional as F
 from torch import nn, einsum, tensor, Tensor, cat, stack, arange, is_tensor
-from torch.utils._pytree import tree_flatten, tree_unflatten
+from torch.utils._pytree import tree_flatten, tree_unflatten, tree_map
 from torch.nn import Module, ModuleList, ModuleDict
 
 from functools import partial, wraps
@@ -81,6 +81,9 @@ def cast_tuple(val, depth = 1):
 def divisible_by(num, den):
     return (num % den) == 0
 
+def detach_all(obj):
+    return tree_map(lambda t: t.detach() if is_tensor(t) and t.requires_grad else t, obj)
+
 def maybe(fn = None):
     if not exists(fn):
         fn = identity
@@ -157,6 +160,19 @@ def or_reduce(masks):
         head = head | rest
     return head
 
+# cache helpers
+
+def get_cached_kvs(
+    cache: LayerIntermediates
+) -> list[tuple[Tensor, Tensor]]:
+
+    cached_kvs = []
+
+    for attn_intermediate in cache.attn_intermediates:
+        cached_kvs.append(attn_intermediate.cached_kv)
+
+    return cached_kvs
+
 # entropy
 
 def calc_entropy(
@@ -1319,7 +1335,7 @@ class Attention(Module):
         value_dim_head = None,
         dim_out = None,
         add_zero_kv = False,         # same as add_zero_attn in pytorch
-        head_learned_sinks = 0,
+        head_learned_sink = False,
         rotate_num_heads = None,
         data_dependent_alibi = False,
         data_dependent_alibi_per_row = False,
@@ -1516,7 +1532,7 @@ class Attention(Module):
             selective = selective,
             custom_attn_fn = custom_attn_fn,
             add_zero_kv = add_zero_kv,
-            head_learned_sinks = head_learned_sinks,
+            head_learned_sink = head_learned_sink,
             flash = flash,
             softclamp_logits = softclamp_logits,
             logit_softclamp_value = logit_softclamp_value,
@@ -2441,8 +2457,13 @@ class AttentionLayers(Module):
         context_pos = None,
         attn_bias = None,
         deep_embeds_and_ids: tuple[nn.Parameter, Tensor] | None = None,
-        self_attn_additional_kv: list[tuple[Tensor, Tensor]] | None = None,
+        self_attn_additional_kv: (
+            LayerIntermediates |
+            list[tuple[Tensor, Tensor]]
+            | None
+        ) = None,
         additional_kv_mask = None,
+        detach_additional_kv = False,
         route_additional_kv_to_top = True,
         condition = None,
         in_attn_cond = None, # https://arxiv.org/abs/2105.04090
@@ -2590,6 +2611,13 @@ class AttentionLayers(Module):
         # additional self attn key / values - say coming from vlm
 
         if exists(self_attn_additional_kv) and route_additional_kv_to_top:
+
+            if isinstance(self_attn_additional_kv, LayerIntermediates):
+                self_attn_additional_kv = get_cached_kvs(self_attn_additional_kv)
+
+            if detach_additional_kv:
+                self_attn_additional_kv = detach_all(self_attn_additional_kv)
+
             num_self_attns = sum([layer_type == 'a' for layer_type in first(layer_variables)])
 
             self_attn_additional_kv = self_attn_additional_kv[-num_self_attns:]