x-transformers 1.39.3__tar.gz → 1.39.4__tar.gz

{x_transformers-1.39.3/x_transformers.egg-info → x_transformers-1.39.4}/PKG-INFO

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

{x_transformers-1.39.3 → x_transformers-1.39.4}/setup.py

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

{x_transformers-1.39.3 → x_transformers-1.39.4}/x_transformers/attend.py

@@ -107,9 +107,9 @@ def qk_l2_dist_squared(q, k):
     l2_dist_squared = torch.cdist(q, k) ** 2
     return unpack_one(l2_dist_squared, packed_shape, '* i j')
 
-# gumbel softmax
+# one-hot straight through softmax
 
-def gumbel_softmax(t, temperature = 1.):
+def one_hot_straight_through(t, temperature = 1.):
     one_hot_indices = t.argmax(dim = -1, keepdim = True)
     one_hot = torch.zeros_like(t).scatter(-1, one_hot_indices, 1.)
 
@@ -180,7 +180,7 @@ class Attend(Module):
         elif sigmoid:
             self.attn_fn = F.sigmoid
         elif hard:
-            self.attn_fn = gumbel_softmax
+            self.attn_fn = one_hot_straight_through
         else:
             softmax_fn = partial(F.softmax, dim = -1)
             self.attn_fn = partial(softmax_fn, dtype = torch.float32) if not qk_norm else softmax_fn

{x_transformers-1.39.3 → x_transformers-1.39.4/x_transformers.egg-info}/PKG-INFO

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