x-transformers 1.32.10__tar.gz → 1.32.12__tar.gz

{x_transformers-1.32.10/x_transformers.egg-info → x_transformers-1.32.12}/PKG-INFO

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

{x_transformers-1.32.10 → x_transformers-1.32.12}/setup.py

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

{x_transformers-1.32.10 → x_transformers-1.32.12}/tests/test_x_transformers.py

@@ -159,7 +159,6 @@ def test_multiple_input_embeds():
     assert embed.shape == (2, 1024, 128)
 
 def test_average_pool_embed():
-
     model = TransformerWrapper(
         num_tokens = 20000,
         max_seq_len = 1024,
@@ -180,7 +179,6 @@ def test_average_pool_embed():
     assert logits.shape == (2, 20000)
 
 def test_cls_token():
-
     model = TransformerWrapper(
         num_tokens = 20000,
         max_seq_len = 1024,
@@ -199,3 +197,24 @@ def test_cls_token():
     logits = model(x, mask = mask)
 
     assert logits.shape == (2, 20000)
+
+def test_squeeze_logit_dim_one():
+    model = TransformerWrapper(
+        num_tokens = 20000,
+        max_seq_len = 1024,
+        logits_dim = 1,
+        average_pool_embed = True,
+        squeeze_out_last_dim = True,
+        attn_layers = Encoder(
+            dim = 128,
+            depth = 6,
+            heads = 8
+        )
+    )
+
+    x = torch.randint(0, 20000, (2, 1024))
+    mask = torch.randint(0, 2, (2, 1024)).bool()
+
+    logits = model(x, mask = mask)
+
+    assert logits.shape == (2,)

{x_transformers-1.32.10 → x_transformers-1.32.12}/x_transformers/x_transformers.py

@@ -45,15 +45,18 @@ def default(val, d):
         return val
     return d() if callable(d) else d
 
-def cast_tuple(val, depth):
+def first(it):
+    return it[0]
+
+def is_empty(x):
+    return len(x) == 0
+
+def cast_tuple(val, depth = 1):
     return val if isinstance(val, tuple) else (val,) * depth
 
 def divisible_by(num, den):
     return (num % den) == 0
 
-def is_empty(x):
-    return len(x) == 0
-
 def maybe(fn):
     @wraps(fn)
     def inner(x, *args, **kwargs):
@@ -1922,6 +1925,7 @@ class TransformerWrapper(Module):
         attn_z_loss_weight = 1e-4,
         average_pool_embed = False,
         use_cls_token = False,
+        squeeze_out_last_dim = False
     ):
         super().__init__()
 
@@ -2006,6 +2010,10 @@ class TransformerWrapper(Module):
 
         self.memory_tokens_interspersed_every = memory_tokens_interspersed_every
 
+        # squeeze out last dimension if possible
+
+        self.squeeze_out_last_dim = squeeze_out_last_dim
+
         # whether can do cached kv decoding
 
         self.can_cache_kv = self.num_memory_tokens == 0
@@ -2173,6 +2181,14 @@ class TransformerWrapper(Module):
             else:
                 logits = self.to_logits(x)
 
+        # maybe squeeze out last dimension of logits
+
+        if self.squeeze_out_last_dim:
+            logits = tuple((rearrange(t, '... 1 -> ...') if t.shape[-1] == 1 else t) for t in cast_tuple(logits))
+
+            if not self.has_multiple_heads:
+                logits = first(logits)
+
         # different returns
 
         if return_logits_and_embeddings:

{x_transformers-1.32.10 → x_transformers-1.32.12/x_transformers.egg-info}/PKG-INFO

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