x-transformers 2.9.0__tar.gz → 2.9.2__tar.gz

{x_transformers-2.9.0 → x_transformers-2.9.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.9.0
+Version: 2.9.2
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers
@@ -2574,4 +2574,16 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@misc{yan2017hierarchicalmultiscaleattentionnetworks,
+    title   = {Hierarchical Multi-scale Attention Networks for Action Recognition}, 
+    author  = {Shiyang Yan and Jeremy S. Smith and Wenjin Lu and Bailing Zhang},
+    year    = {2017},
+    eprint  = {1708.07590},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CV},
+    url     = {https://arxiv.org/abs/1708.07590}, 
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-2.9.0 → x_transformers-2.9.2}/README.md

@@ -2525,4 +2525,16 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@misc{yan2017hierarchicalmultiscaleattentionnetworks,
+    title   = {Hierarchical Multi-scale Attention Networks for Action Recognition}, 
+    author  = {Shiyang Yan and Jeremy S. Smith and Wenjin Lu and Bailing Zhang},
+    year    = {2017},
+    eprint  = {1708.07590},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CV},
+    url     = {https://arxiv.org/abs/1708.07590}, 
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-2.9.0 → x_transformers-2.9.2}/pyproject.toml

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

{x_transformers-2.9.0 → x_transformers-2.9.2}/tests/test_x_transformers.py

@@ -1373,3 +1373,15 @@ def test_muon_params():
 
     enc = Encoder(dim = 512, depth = 2)
     assert len(enc.muon_parameters()) > 0
+
+def test_stochastic_attn():
+    from x_transformers import Attention
+
+    attn = Attention(dim = 512, gumbel_softmax = True)
+    out, intermediate = attn(torch.randn(1, 1024, 512), return_intermediates = True)
+
+    assert out.shape == (1, 1024, 512)
+
+    from x_transformers.attend import log_prob_from_hard_attend
+    log_probs = log_prob_from_hard_attend(intermediate)
+    assert log_probs.shape == (1, 8, 1024)

{x_transformers-2.9.0 → x_transformers-2.9.2}/x_transformers/attend.py

@@ -67,6 +67,15 @@ def once(fn):
 
 print_once = once(print)
 
+# gumbel softmax attention related
+
+def log_prob_from_hard_attend(intermeds: Intermediates):
+    log_probs = intermeds.pre_softmax_attn.log_softmax(dim = -1)
+
+    one_hot = intermeds.post_softmax_attn.argmax(dim = -1, keepdim = True)
+    log_prob = log_probs.gather(-1, one_hot)
+    return rearrange(log_prob, 'b h i 1 -> b h i')
+
 # selective attention
 # https://arxiv.org/abs/2410.02703 - section 3.3
 # it is a technique to allow each token to prevent itself from being attended to by future tokens
@@ -171,6 +180,9 @@ class Attend(Module):
         qk_norm = False,
         l2_distance = False,
         sigmoid = False,
+        gumbel_softmax = False,
+        gumbel_softmax_temp = 1.,
+        gumbel_softmax_hard = True,
         custom_attn_fn: Callable | None = None,
         flash = False,
         softclamp_logits = False,
@@ -203,7 +215,7 @@ class Attend(Module):
         assert not (flash and hard), 'hard attention not available for flash'
         assert not (flash and is_sparse_topk_attn), 'topk attention not available for flash'
 
-        assert at_most_one_of(sigmoid, hard, l2_distance, is_sparse_topk_attn)
+        assert at_most_one_of(sigmoid, hard, l2_distance, gumbel_softmax, is_sparse_topk_attn)
 
         if exists(custom_attn_fn):
             self.attn_fn = custom_attn_fn
@@ -213,6 +225,8 @@ class Attend(Module):
             self.attn_fn = one_hot_straight_through
         elif is_sparse_topk_attn:
             self.attn_fn = partial(sparse_topk_attn, sparse_topk = sparse_topk, straight_through = sparse_topk_straight_through)
+        elif gumbel_softmax:
+            self.attn_fn = partial(F.gumbel_softmax, dim = -1, tau = gumbel_softmax_temp, hard = gumbel_softmax_hard)
         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-2.9.0 → x_transformers-2.9.2}/x_transformers/x_transformers.py

@@ -1336,6 +1336,9 @@ class Attention(Module):
         value_rmsnorm = False,      # used in alphagenome and bytedance's GR3 for further stability
         l2_distance = False,
         sigmoid = False,
+        gumbel_softmax = False,
+        gumbel_softmax_temp = 1.,
+        gumbel_softmax_hard = True,
         selective = False,
         custom_attn_fn: Callable | None = None,
         hybrid_module: Module | None = None,
@@ -1541,6 +1544,9 @@ class Attention(Module):
             scale = qk_norm_scale if qk_norm else self.scale,
             l2_distance = l2_distance,
             sigmoid = sigmoid,
+            gumbel_softmax = gumbel_softmax,
+            gumbel_softmax_temp = gumbel_softmax_temp,
+            gumbel_softmax_hard = gumbel_softmax_hard,
             selective = selective,
             custom_attn_fn = custom_attn_fn,
             add_zero_kv = add_zero_kv,