x-transformers 1.38.3__tar.gz → 1.39.0__tar.gz

{x_transformers-1.38.3/x_transformers.egg-info → x_transformers-1.39.0}/PKG-INFO

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

{x_transformers-1.38.3 → x_transformers-1.39.0}/README.md

@@ -549,7 +549,8 @@ model = TransformerWrapper(
         dim = 512,
         depth = 6,
         heads = 8,
-        attn_talking_heads = True  # turn on information exchange between attention heads
+        attn_pre_talking_heads = True,  # linear combination across pre-softmax attn logits across heads
+        attn_post_talking_heads = True  # linear combination across post-softmax attn across heads
     )
 )
 ```

{x_transformers-1.38.3 → x_transformers-1.39.0}/setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'x-transformers',
   packages = find_packages(exclude=['examples']),
-  version = '1.38.3',
+  version = '1.39.0',
   license='MIT',
   description = 'X-Transformers - Pytorch',
   author = 'Phil Wang',

{x_transformers-1.38.3 → x_transformers-1.39.0}/x_transformers/attend.py

@@ -128,7 +128,8 @@ class Attend(Module):
         dropout = 0.,
         causal = False,
         heads = None,
-        talking_heads = False,
+        pre_talking_heads = True,
+        post_talking_heads = True,
         sparse_topk = None,
         scale = None,
         qk_norm = False,
@@ -179,14 +180,15 @@ class Attend(Module):
 
         # talking heads
 
-        assert not (flash and talking_heads), 'talking heads not compatible with flash attention'
+        assert not (flash and (pre_talking_heads or post_talking_heads)), 'talking heads not compatible with flash attention'
 
-        self.talking_heads = talking_heads
-        if talking_heads:
-            self.pre_softmax_talking_heads = nn.Conv2d(heads, heads, 1, bias = False)
-            self.post_softmax_talking_heads = nn.Conv2d(heads, heads, 1, bias = False)
+        self.pre_softmax_talking_heads = nn.Conv2d(heads, heads, 1, bias = False) if pre_talking_heads else None
+        self.post_softmax_talking_heads = nn.Conv2d(heads, heads, 1, bias = False) if post_talking_heads else None
 
+        if exists(self.pre_softmax_talking_heads):
             nn.init.dirac_(self.pre_softmax_talking_heads.weight)
+
+        if exists(self.post_softmax_talking_heads):
             nn.init.dirac_(self.post_softmax_talking_heads.weight)
 
         # selective attention
@@ -434,8 +436,8 @@ class Attend(Module):
 
         qk_similarities = sim.clone()
 
-        if self.talking_heads:
-            sim = self.pre_softmax_talking_heads(sim)
+        if exists(self.pre_softmax_talking_heads):
+            sim = sim + self.pre_softmax_talking_heads(sim)
 
         if exists(attn_bias):
             sim = sim + attn_bias
@@ -482,7 +484,7 @@ class Attend(Module):
 
         attn = self.attn_dropout(attn)
 
-        if self.talking_heads:
+        if exists(self.post_softmax_talking_heads):
             attn = self.post_softmax_talking_heads(attn)
 
         out = einsum(f'b h i j, {kv_einsum_eq} -> b h i d', attn, v)

{x_transformers-1.38.3 → x_transformers-1.39.0}/x_transformers/x_transformers.py

@@ -907,7 +907,8 @@ class Attention(Module):
         heads = 8,
         causal = False,
         flash = False,
-        talking_heads = False,
+        pre_talking_heads = False,
+        post_talking_heads = False,
         head_scale = False,
         sparse_topk = None,
         num_mem_kv = 0,
@@ -1036,7 +1037,8 @@ class Attention(Module):
         self.attend = Attend(
             heads = heads,
             causal = causal,
-            talking_heads = talking_heads,
+            pre_talking_heads = pre_talking_heads,
+            post_talking_heads = post_talking_heads,
             dropout = dropout,
             sparse_topk = sparse_topk,
             qk_norm = qk_norm,

{x_transformers-1.38.3 → x_transformers-1.39.0/x_transformers.egg-info}/PKG-INFO

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