softmax is already stable, also start enforcing it to be always float32

NopenAI · Aug 7, 2022 · 1adb2ea · 1adb2ea
1 parent 7212069
commit 1adb2ea
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 6 deletions.
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'x-transformers',
   packages = find_packages(exclude=['examples']),
-  version = '0.32.2',
+  version = '0.32.3',
   license='MIT',
   description = 'X-Transformers - Pytorch',
   author = 'Phil Wang',

diff --git a/x_transformers/x_transformers.py b/x_transformers/x_transformers.py
@@ -66,10 +66,6 @@ def l2norm(t, groups = 1):
     t = F.normalize(t, p = 2, dim = -1)
     return rearrange(t, '... g d -> ... (g d)')
 
-def stable_softmax(t, dim = -1):
-    t = t - t.amax(dim = dim, keepdim = True).detach()
-    return F.softmax(t, dim = dim)
-
 # init helpers
 
 def init_zero_(layer):
@@ -573,7 +569,7 @@ def __init__(
         self.sparse_topk = sparse_topk
 
         # entmax
-        self.attn_fn = entmax15 if use_entmax15 else stable_softmax
+        self.attn_fn = entmax15 if use_entmax15 else partial(F.softmax, dtype = torch.float32)
 
         # add memory key / values
         self.num_mem_kv = num_mem_kv