x-transformers 2.8.0__tar.gz → 2.8.1__tar.gz

{x_transformers-2.8.0 → x_transformers-2.8.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.8.0
+Version: 2.8.1
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers

{x_transformers-2.8.0 → x_transformers-2.8.1}/pyproject.toml

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

{x_transformers-2.8.0 → x_transformers-2.8.1}/tests/test_x_transformers.py

@@ -1340,3 +1340,23 @@ def test_qk_clip_attn_layers():
     out, intermediates = model(seq, return_intermediates = True)
 
     model.attn_qk_clip_(intermediates)
+
+def test_vae():
+    from x_transformers.gpt_vae import GPTVAE
+
+    model = GPTVAE(
+        num_tokens = 256,
+        max_seq_len = 1024,
+        dim = 512,
+        depth = 4,
+        enc_depth = 2
+    )
+
+    seq = torch.randint(0, 256, (1, 1024))
+
+    loss = model(seq)
+    loss.backward()
+
+    style = torch.randint(0, 256, (1, 1024))
+
+    out = model.generate(seq[:, :512], 512, seq_for_latents = style)

{x_transformers-2.8.0 → x_transformers-2.8.1}/x_transformers/gpt_vae.py

@@ -68,10 +68,13 @@ class GPTVAE(Module):
 
         self.to_latent_mean_log_variance = nn.Sequential(
             nn.Linear(dim, dim_latent * 2),
-            Rearrange('b (two d) -> two b 1 d', two = 2)
+            Rearrange('b (two d) -> two b d', two = 2)
         )
 
-        self.from_latent_to_prepend_token = nn.Linear(dim_latent, dim)
+        self.from_latent_to_prepend_token = nn.Sequential(
+            nn.Linear(dim_latent, dim),
+            Rearrange('b d -> b 1 d')
+        )
 
         self.decoder = TransformerWrapper(
             num_tokens = num_tokens,
@@ -126,11 +129,19 @@ class GPTVAE(Module):
         prompts,
         seq_len,
         latents = None,
+        seq_for_latents = None,
         **generate_kwargs
     ):
         assert prompts.ndim in {1, 2}
         batch = prompts.shape[0] if prompts.ndim == 2 else 1
 
+        # if seq_for_latents passed in, derive latents from it
+
+        if exists(seq_for_latents):
+            assert not exists(latents), 'latents should not be passed in if given the seq from which to derive them'
+
+            latents = self.encode_to_latents(seq_for_latents)
+
         # prepend embeds
 
         prepend_embeds = None
@@ -143,9 +154,6 @@ class GPTVAE(Module):
 
             prepend_embeds = self.from_latent_to_prepend_token(latents)
 
-        if exists(prepend_embeds):
-            prepend_embeds = rearrange(prepend_embeds, 'b d -> b 1 d')
-
         # generated
 
         generated = self.ar_wrapped_decoder.generate(