x-transformers 1.40.9__tar.gz → 1.40.11__tar.gz

{x_transformers-1.40.9/x_transformers.egg-info → x_transformers-1.40.11}/PKG-INFO

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

{x_transformers-1.40.9 → x_transformers-1.40.11}/setup.py

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

{x_transformers-1.40.9 → x_transformers-1.40.11}/tests/test_x_transformers.py

@@ -179,12 +179,14 @@ def test_average_pool_embed():
 
     assert logits.shape == (2, 20000)
 
-def test_cls_token():
+@pytest.mark.parametrize('num_cls_tokens', (1, 2))
+def test_cls_token(num_cls_tokens):
     model = TransformerWrapper(
         num_tokens = 20000,
         max_seq_len = 1024,
         num_memory_tokens = 2,
         use_cls_token = True,
+        num_cls_tokens=num_cls_tokens,
         attn_layers = Encoder(
             dim = 128,
             depth = 6,
@@ -197,7 +199,12 @@ def test_cls_token():
 
     logits = model(x, mask = mask)
 
-    assert logits.shape == (2, 20000)
+    if num_cls_tokens == 1:
+        expected_shape = (2, 20000)
+    else:
+        expected_shape = (2, num_cls_tokens, 20000)
+
+    assert logits.shape == expected_shape
 
 def test_squeeze_logit_dim_one():
     model = TransformerWrapper(

{x_transformers-1.40.9 → x_transformers-1.40.11}/x_transformers/x_transformers.py

@@ -1019,7 +1019,7 @@ class Attention(Module):
         self.qk_norm_q_scale = self.qk_norm_k_scale = 1
         if qk_norm and qk_norm_dim_scale:
             self.qk_norm_q_scale = nn.Parameter(torch.ones(heads, 1, dim_head))
-            self.qk_norm_k_scale = nn.Parameter(torch.ones(heads, 1, dim_head))
+            self.qk_norm_k_scale = nn.Parameter(torch.ones(kv_heads, 1, dim_head))
 
         assert (not qk_norm) or divisible_by(dim_head, qk_norm_groups), 'dimension per attention head must be divisible by the qk norm groups'
         assert not (qk_norm and (dim_head // qk_norm_groups) <= 2), 'the group dimension may be too small (2 was too small in my tests, but 4 still works, surprisingly)'
@@ -2104,6 +2104,7 @@ class TransformerWrapper(Module):
         attn_z_loss_weight = 1e-4,
         average_pool_embed = False,
         use_cls_token = False,
+        num_cls_tokens = 1,
         squeeze_out_last_dim = False,
         token_emb: TokenEmbedding | None = None,
         mixture_of_softmax = False,
@@ -2116,6 +2117,7 @@ class TransformerWrapper(Module):
         emb_dim = default(emb_dim, dim)
         self.emb_dim = emb_dim
         self.num_tokens = num_tokens
+        self.num_cls_tokens = num_cls_tokens
 
         self.max_seq_len = max_seq_len
         self.max_mem_len = max_mem_len
@@ -2172,7 +2174,7 @@ class TransformerWrapper(Module):
         self.cls_token = None
 
         if use_cls_token:
-            self.cls_token = nn.Parameter(torch.zeros(dim))
+            self.cls_token = nn.Parameter(torch.zeros(num_cls_tokens, dim))
             nn.init.normal_(self.cls_token, std = 0.02)
 
         # whether to average pool the embed (`global average pool`)
@@ -2329,11 +2331,11 @@ class TransformerWrapper(Module):
         # maybe cls token
 
         if exists(self.cls_token):
-            cls_tokens = repeat(self.cls_token, 'd -> b d', b = b)
+            cls_tokens = repeat(self.cls_token, '... -> b ...', b = b)
             x, cls_packed_shape = pack([cls_tokens, x], 'b * d')
 
             if exists(mask):
-                mask = F.pad(mask, (1, 0), value = True)
+                mask = F.pad(mask, (self.num_cls_tokens, 0), value = True)
 
         # maybe memory / register tokens
 
@@ -2415,6 +2417,7 @@ class TransformerWrapper(Module):
 
         if exists(self.cls_token):
             x, _ = unpack(x, cls_packed_shape, 'b * d')
+            x = x.squeeze(1)  # Remove sequence dimension if num_cls_tokens=1 to keep previous behavior
 
         # handle expansion to mixture if needed (for mixture of softmax)
 

{x_transformers-1.40.9 → x_transformers-1.40.11/x_transformers.egg-info}/PKG-INFO

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