x-transformers 2.11.12__tar.gz → 2.11.14__tar.gz

{x_transformers-2.11.12 → x_transformers-2.11.14}/PKG-INFO

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

{x_transformers-2.11.12 → x_transformers-2.11.14}/pyproject.toml

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

{x_transformers-2.11.12 → x_transformers-2.11.14}/tests/test_x_transformers.py

@@ -1437,6 +1437,11 @@ def test_free(
 
     assert aux_loss.numel() == 1
 
+    rand_indices = torch.randint(0, 2 ** 8, ())
+    generated = model.generate(seq[:, :1], 32, latents = rand_indices)
+
+    assert generated.shape == (1, 32)
+
 def test_kv_input_residual():
     attn = Decoder(
         dim = 256,

{x_transformers-2.11.12 → x_transformers-2.11.14}/train_free.py

@@ -63,8 +63,7 @@ model = FreeTransformer(
     latent_bits = LATENT_BITS
 ).cuda()
 
-rand_index = torch.randint(0, 2 ** LATENT_BITS, ())
-latents = F.one_hot(rand_index, 2 ** LATENT_BITS).float().cuda()
+one_hot_indices = torch.randint(0, 2 ** LATENT_BITS, ())
 
 # prepare enwik8 data
 
@@ -126,9 +125,9 @@ for i in tqdm.tqdm(range(NUM_BATCHES), mininterval=10., desc='training'):
         sample = model.generate(
             prompts = inp,
             seq_len = GENERATE_LENGTH,
-            latents = latents
+            latents = one_hot_indices
         )
 
         output_str = decode_tokens(sample)
 
-        print(f'\n\nlatent {rand_index.tolist()} - ', output_str)
+        print(f'\n\nlatent {one_hot_indices.tolist()} - ', output_str)

{x_transformers-2.11.12 → x_transformers-2.11.14}/x_transformers/free_transformer.py

@@ -282,6 +282,10 @@ class FreeTransformer(Module):
             if not is_tensor(latents):
                 latents = tensor(latents, device = self.device)
 
+            if latents.dtype in (torch.int, torch.long):
+                # if given as indices
+                latents = F.one_hot(latents, self.binary_mapper.num_codes).float()
+
             if latents.ndim == 1: # repeat latents
                 latents = repeat(latents, 'd -> b 1 d', b = batch)
             elif latents.ndim == 2: