x-transformers 1.31.4__tar.gz → 1.31.5__tar.gz

{x_transformers-1.31.4/x_transformers.egg-info → x_transformers-1.31.5}/PKG-INFO

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

{x_transformers-1.31.4 → x_transformers-1.31.5}/setup.py

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

{x_transformers-1.31.4 → x_transformers-1.31.5}/tests/test_x_transformers.py

@@ -1,5 +1,4 @@
 import torch
-
 from x_transformers.x_transformers import (
     XTransformer,
     TransformerWrapper,
@@ -111,3 +110,20 @@ def test_adaptive_rmsnorm():
     condition = torch.randn(2, 768)
 
     model(x, condition = condition)
+
+def test_attn_softclamp_logits():
+    model = TransformerWrapper(
+        num_tokens = 20000,
+        max_seq_len = 1024,
+        attn_layers = Decoder(
+            dim = 512,
+            dim_condition = 768,
+            depth = 12,
+            heads = 8,
+            attn_softclamp_logits = True,
+        )
+    )
+
+    x = torch.randint(0, 256, (1, 1024))
+
+    model(x)

{x_transformers-1.31.4 → x_transformers-1.31.5}/x_transformers/attend.py

@@ -69,7 +69,7 @@ def onnx_create_causal_mask(i, j, device):
 
 # main class
 
-class Attend(nn.Module):
+class Attend(Module):
     def __init__(
         self,
         *,
@@ -81,7 +81,8 @@ class Attend(nn.Module):
         scale = None,
         qk_norm = False,
         flash = False,
-        logit_softclamp_value = None,
+        softclamp_logits = False,
+        logit_softclamp_value = 30.,
         add_zero_kv = False,
         cope = None,
         onnxable = False,
@@ -123,10 +124,11 @@ class Attend(nn.Module):
 
         # soft clamp attention logit value
 
-        if exists(logit_softclamp_value):
+        if softclamp_logits:
             assert not flash, 'flash attention not compatible with logit softclamp value yet'
             assert logit_softclamp_value > 0.
 
+        self.softclamp_logits = softclamp_logits
         self.logit_softclamp_value = logit_softclamp_value
 
         # contextual positional encoding
@@ -308,6 +310,9 @@ class Attend(nn.Module):
         if exists(attn_bias):
             sim = sim + attn_bias
 
+        if self.softclamp_logits:
+            sim = softclamp(sim, self.logit_softclamp_value)
+
         i, j, dtype = *sim.shape[-2:], sim.dtype
 
         mask_value = -torch.finfo(sim.dtype).max
@@ -329,9 +334,6 @@ class Attend(nn.Module):
 
         pre_softmax_attn = sim.clone()
 
-        if exists(self.logit_softclamp_value):
-            sim = softclamp(sim, self.logit_softclamp_value)
-
         attn = self.attn_fn(sim, dim = -1)
         attn = attn.type(dtype)
 

{x_transformers-1.31.4 → x_transformers-1.31.5}/x_transformers/x_transformers.py

@@ -884,7 +884,8 @@ class Attention(Module):
         cope_max_pos = 16,
         cope_soft_onehot_pos = False,
         cope_talking_heads = False,
-        logit_softclamp_value = None,
+        softclamp_logits = False,
+        logit_softclamp_value = 30.,
         onnxable = False
     ):
         super().__init__()
@@ -987,6 +988,7 @@ class Attention(Module):
             scale = qk_norm_scale if qk_norm else self.scale,
             add_zero_kv = add_zero_kv,
             flash = flash,
+            softclamp_logits = softclamp_logits,
             logit_softclamp_value = logit_softclamp_value,
             cope = cope,
             onnxable = onnxable

{x_transformers-1.31.4 → x_transformers-1.31.5/x_transformers.egg-info}/PKG-INFO

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