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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.11.11
+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.11 → x_transformers-2.11.14}/pyproject.toml

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

{x_transformers-2.11.11 → 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.11 → 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.11 → 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:

{x_transformers-2.11.11 → x_transformers-2.11.14}/x_transformers/x_transformers.py

@@ -740,11 +740,14 @@ def apply_rotary_pos_emb(t, freqs, scale = 1):
     rot_dim, seq_len, orig_dtype = freqs.shape[-1], t.shape[-2], t.dtype
 
     freqs = freqs[:, -seq_len:, :]
-    scale = scale[:, -seq_len:, :] if isinstance(scale, torch.Tensor) else scale
+    scale = scale[:, -seq_len:, :] if is_tensor(scale) else scale
 
     if t.ndim == 4 and freqs.ndim == 3:
         freqs = rearrange(freqs, 'b n d -> b 1 n d')
 
+        if is_tensor(scale):
+            scale = rearrange(scale, 'b n d -> b 1 n d')
+
     # partial rotary embeddings, Wang et al. GPT-J
     t, t_unrotated = t[..., :rot_dim], t[..., rot_dim:]
     t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale)