final touch-up to free transformer for the day

Author: lucidrains

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" }

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,

train_free.py

@@ -63,8 +63,7 @@ def decode_tokens(tokens):
     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 @@ def __len__(self):
         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/free_transformer.py

@@ -282,6 +282,10 @@ def generate(
             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: