x-transformers 2.11.24__tar.gz → 2.12.1__tar.gz

{x_transformers-2.11.24 → x_transformers-2.12.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.11.24
+Version: 2.12.1
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers
@@ -2630,4 +2630,16 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@misc{gopalakrishnan2025decouplingwhatwherepolar,
+    title   = {Decoupling the "What" and "Where" With Polar Coordinate Positional Embeddings}, 
+    author  = {Anand Gopalakrishnan and Robert Csordás and Jürgen Schmidhuber and Michael C. Mozer},
+    year    = {2025},
+    eprint  = {2509.10534},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2509.10534}, 
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-2.11.24 → x_transformers-2.12.1}/README.md

@@ -2581,4 +2581,16 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@misc{gopalakrishnan2025decouplingwhatwherepolar,
+    title   = {Decoupling the "What" and "Where" With Polar Coordinate Positional Embeddings}, 
+    author  = {Anand Gopalakrishnan and Robert Csordás and Jürgen Schmidhuber and Michael C. Mozer},
+    year    = {2025},
+    eprint  = {2509.10534},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2509.10534}, 
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-2.11.24 → x_transformers-2.12.1}/pyproject.toml

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

{x_transformers-2.11.24 → x_transformers-2.12.1}/tests/test_x_transformers.py

@@ -1508,3 +1508,21 @@ def test_derf():
     x = torch.randint(0, 256, (1, 10))
 
     logits = model(x)
+
+def test_pope():
+    from x_transformers import TransformerWrapper, Decoder
+
+    model = TransformerWrapper(
+        num_tokens = 256,
+        max_seq_len = 1024,
+        attn_layers = Decoder(
+            dim = 512,
+            depth = 6,
+            heads = 8,
+            polar_pos_emb = True,
+        )
+    )
+
+    x = torch.randint(0, 256, (1, 10))
+
+    logits = model(x)

{x_transformers-2.11.24 → x_transformers-2.12.1}/train_length_extrapolate.py

@@ -1,3 +1,11 @@
+# /// script
+# dependencies = [
+#   "accelerate",
+#   "tqdm",
+#   "x-transformers>=2.12.0",
+# ]
+# ///
+
 from x_transformers import TransformerWrapper, Decoder
 from x_transformers.autoregressive_wrapper import AutoregressiveWrapper
 
@@ -10,6 +18,8 @@ import torch.optim as optim
 from torch.nn import functional as F
 from torch.utils.data import DataLoader, Dataset
 
+from accelerate import Accelerator
+
 # constants
 
 NUM_BATCHES = int(1e5)
@@ -20,7 +30,7 @@ GENERATE_EVERY  = 500
 GENERATE_LENGTH = 256
 SEQ_LEN = 256
 
-VALIDATE_EVERY  = 100
+VALIDATE_EVERY  = 250
 VALIDATE_SEQ_LENS = (256, 512, 1024, 2048, 4096)
 
 # helpers
@@ -36,6 +46,10 @@ def decode_token(token):
 def decode_tokens(tokens):
     return ''.join(list(map(decode_token, tokens)))
 
+# accelerator
+
+accelerator = Accelerator()
+
 # instantiate GPT-like decoder model
 
 model = TransformerWrapper(
@@ -46,12 +60,13 @@ model = TransformerWrapper(
         dim = 512,
         depth = 6,
         heads = 8,
-        dynamic_pos_bias = True,
+        polar_pos_emb = True,
+        rotary_pos_emb = False,
+        dynamic_pos_bias = False
     )
 )
 
 model = AutoregressiveWrapper(model)
-model.cuda()
 
 # prepare enwik8 data
 
@@ -69,69 +84,84 @@ class TextSamplerDataset(Dataset):
     def __getitem__(self, index):
         rand_start = torch.randint(0, self.data.size(0) - self.seq_len - 1, (1,))
         full_seq = self.data[rand_start: rand_start + self.seq_len + 1].long()
-        return full_seq.cuda()
+        return full_seq
 
     def __len__(self):
         return self.data.size(0) // self.seq_len
 
 train_dataset = TextSamplerDataset(data_train, SEQ_LEN)
-train_loader  = cycle(DataLoader(train_dataset, batch_size = BATCH_SIZE, drop_last = True))
+train_loader  = DataLoader(train_dataset, batch_size = BATCH_SIZE, drop_last = True)
 
 val_dataset_generate = TextSamplerDataset(data_val, SEQ_LEN)
 
+# optimizer
+
+optim = torch.optim.Adam(model.parameters(), lr = LEARNING_RATE)
+
+# prepare
+
+model, optim, train_loader = accelerator.prepare(model, optim, train_loader)
+
+train_loader = cycle(train_loader)
+
 # validation loaders with different sequence lengths
 
 val_loaders = dict()
 
 for valid_seq_len in VALIDATE_SEQ_LENS:
     val_dataset   = TextSamplerDataset(data_val, valid_seq_len)
-    val_loader    = cycle(DataLoader(val_dataset, batch_size = BATCH_SIZE, drop_last = True))
+    val_loader    = DataLoader(val_dataset, batch_size = BATCH_SIZE, drop_last = True)
+    val_loader    = cycle(val_loader)
 
     val_loaders[valid_seq_len] = val_loader
 
-# optimizer
-
-optim = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
-
 # training
 
 for i in tqdm.tqdm(range(NUM_BATCHES), mininterval=10., desc='training'):
     model.train()
 
-    for __ in range(GRADIENT_ACCUMULATE_EVERY):
-        loss = model(next(train_loader))
-        (loss / GRADIENT_ACCUMULATE_EVERY).backward()
+    for _ in range(GRADIENT_ACCUMULATE_EVERY):
+        data = next(train_loader)
+        loss = model(data)
+        accelerator.backward(loss / GRADIENT_ACCUMULATE_EVERY)
+
+        if accelerator.sync_gradients:
+            accelerator.clip_grad_norm_(model.parameters(), 0.5)
 
-    print(f'training loss: {loss.item()}')
+        optim.step()
+        optim.zero_grad()
 
-    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
-    optim.step()
-    optim.zero_grad()
+    if i % 10 == 0:
+        accelerator.print(f'training loss: {loss.item()}')
 
     if i % VALIDATE_EVERY == 0:
-        print(f'validation losses:\n')
+        accelerator.print(f'validation losses:\n')
 
         model.eval()
-        with torch.no_grad():
+        with torch.inference_mode():
             for valid_seq_len in VALIDATE_SEQ_LENS:
                 val_loader = val_loaders[valid_seq_len]
 
-                loss = model(next(val_loader))
-                print(f'[{valid_seq_len}]:\t {loss.item()}')
+                val_data = next(val_loader).to(accelerator.device)
+                loss = model(val_data)
+                accelerator.print(f'[{valid_seq_len}]:\t {loss.item()}')
 
-        print('\n')
+        accelerator.print('\n')
 
     if i % GENERATE_EVERY == 0:
         model.eval()
+        unwrapped_model = accelerator.unwrap_model(model)
+        
         inp = random.choice(val_dataset_generate)[:-1]
+        inp = inp.to(accelerator.device)
         prime = decode_tokens(inp)
-        print(f'%s \n\n %s', (prime, '*' * 100))
+        accelerator.print(f'{prime} \n\n {"*" * 100}')
 
-        sample = model.generate(
+        sample = unwrapped_model.generate(
             prompts = inp,
             seq_len = GENERATE_LENGTH,
             cache_kv = True
         )
 
         output_str = decode_tokens(sample)
-        print(f'{output_str}\n\n')
+        accelerator.print(f'{output_str}\n\n')

{x_transformers-2.11.24 → x_transformers-2.12.1}/x_transformers/x_transformers.py

@@ -779,6 +779,49 @@ def apply_rotary_pos_emb(t, freqs, scale = 1):
 
     return out.type(orig_dtype)
 
+class PolarEmbedding(Module):
+    """ https://arxiv.org/abs/2509.10534 """
+
+    def __init__(
+        self,
+        dim,
+        bias_uniform_init = False,
+        base = 10000,
+    ):
+        super().__init__()
+        inv_freq = 1. / (base ** (arange(0, dim).float() / dim))
+        self.register_buffer('inv_freq', inv_freq)
+
+        self.learned_bias = nn.Parameter(torch.zeros(dim))
+
+        if bias_uniform_init:
+            self.learned_bias.uniform_(-2. * math.pi, 0.)
+
+    @autocast('cuda', enabled = False)
+    def forward(self, t, offset = 0):
+        max_pos = t.max() + 1
+
+        if t.ndim == 1:
+            t = rearrange(t, 'n -> 1 n')
+
+        freqs = torch.einsum('b i , j -> b i j', t.type_as(self.inv_freq), self.inv_freq)
+
+        bias = self.learned_bias.clamp(-2. * math.pi, 0.)
+
+        return freqs, bias
+
+@autocast('cuda', enabled = False)
+def apply_polar_pos_emb(t, freqs):
+    rot_dim, seq_len, orig_dtype = freqs.shape[-1], t.shape[-2], t.dtype
+    freqs = freqs[:, -seq_len:]
+
+    t = t.float()
+
+    t = F.softplus(t)
+    out = cat((t * freqs.cos(), t * freqs.sin()), dim = -1)
+
+    return out.type(orig_dtype)
+
 # norms
 
 class Scale(Module):
@@ -1745,6 +1788,7 @@ class Attention(Module):
         attn_bias = None,
         rotary_pos_emb = None,
         context_rotary_pos_emb = None,
+        polar_pos_emb = None,
         pos = None, # for custom alibi positions
         prev_attn = None,
         mem = None,
@@ -1896,6 +1940,11 @@ class Attention(Module):
                 q = cat((q_rest, q), dim = 1)
                 k = cat((k_rest, k), dim = 1)
 
+        if exists(polar_pos_emb):
+            freqs, bias = polar_pos_emb
+            q = apply_polar_pos_emb(q, freqs)
+            k = apply_polar_pos_emb(k, freqs + bias)
+
         input_mask = context_mask
 
         if not exists(input_mask) and not has_context:
@@ -2174,6 +2223,8 @@ class AttentionLayers(Module):
         rotary_xpos_scale_base = 512,
         rotary_base_rescale_factor = 1.,
         rotate_num_heads = None,
+        polar_pos_emb = False,
+        polar_bias_uniform_init = False,
         weight_tie_layers = False,
         custom_layers: tuple[str, ...] | None = None,
         layers_execute_order: tuple[int, ...] | None = None,
@@ -2250,14 +2301,13 @@ class AttentionLayers(Module):
 
         # LIMe
 
-        hiddens_counter = 0
         self.layer_integrators = ModuleList([])
 
         assert not (qkv_receive_diff_residuals and not (hyper_conn_produce_diff_views or integrate_layers))
 
         # positions related
 
-        self.disable_abs_pos_emb = default(disable_abs_pos_emb, (rel_pos_bias or rotary_pos_emb))
+        self.disable_abs_pos_emb = default(disable_abs_pos_emb, (rel_pos_bias or rotary_pos_emb or polar_pos_emb))
 
         rotary_emb_dim = default(rotary_emb_dim, dim_head // 2)
 
@@ -2266,9 +2316,14 @@ class AttentionLayers(Module):
         if verbose and rotary_emb_dim < 32:
             logger.warning('when training language model, rotary embedding dimension should be at least 32')
 
+        assert at_most_one_of(rotary_pos_emb, polar_pos_emb), f'either rotary positional embedding or polar positional embedding can be turned on'
         assert not (rotary_xpos and not causal), 'rotary xpos is not compatible with bidirectional attention'
         self.rotary_pos_emb = RotaryEmbedding(rotary_emb_dim, use_xpos = rotary_xpos, scale_base = rotary_xpos_scale_base, interpolation_factor = rotary_interpolation_factor, base_rescale_factor = rotary_base_rescale_factor) if rotary_pos_emb else None
 
+        # polar positional embedding (PoPE) - https://arxiv.org/abs/2509.10534
+
+        self.polar_pos_emb = PolarEmbedding(dim_head, polar_bias_uniform_init) if polar_pos_emb else None
+
         assert at_most_one_of(alibi_pos_bias, rel_pos_bias, data_dependent_alibi), 'you can only choose one of Alibi positional bias, data dependent Alibi (forgetting transformers), dynamic tanh, or T5 relative positional bias'
         assert rel_pos_num_buckets <= rel_pos_max_distance, 'number of relative position buckets must be less than the relative position max distance'
 
@@ -2626,6 +2681,7 @@ class AttentionLayers(Module):
         cache_age = 1,
         return_hiddens = False,
         rotary_pos_emb = None,
+        polar_pos_emb = None,
         pos = None,
         context_pos = None,
         attn_bias = None,
@@ -2721,6 +2777,15 @@ class AttentionLayers(Module):
                     context_rotary_pos_emb = context_rotary_pos_emb
                 )
 
+        # polar positions
+
+        if exists(self.polar_pos_emb):
+            if not exists(polar_pos_emb):
+                if not exists(pos):
+                    pos = arange(x.shape[1] + seq_pos_offset, device = x.device)
+
+                polar_pos_emb = self.polar_pos_emb(pos)
+
         # assume cached key / values
 
         prev_cache_length = 0
@@ -2910,7 +2975,7 @@ class AttentionLayers(Module):
             # forward depending on layer type
 
             if layer_type == 'a':
-                out, inter = block(x, mask = mask, context_mask = self_attn_kv_mask, attn_mask = attn_mask, rel_pos = self.rel_pos, pos = pos, rotary_pos_emb = rotary_pos_emb, additional_key_values = next(iter_self_attn_kv, None), additional_key_value_mask = additional_kv_mask, prev_attn = prev_attn, cache = next(iter_attn_cache, None), mem = layer_mem, mem_mask = layer_mem_mask, attn_bias = attn_bias, kv_input_residual = next(self_attn_kv_residuals_iter, None), value_residual = maybe_self_attn_value_residual, return_intermediates = True)
+                out, inter = block(x, mask = mask, context_mask = self_attn_kv_mask, attn_mask = attn_mask, rel_pos = self.rel_pos, pos = pos, rotary_pos_emb = rotary_pos_emb, polar_pos_emb = polar_pos_emb, additional_key_values = next(iter_self_attn_kv, None), additional_key_value_mask = additional_kv_mask, prev_attn = prev_attn, cache = next(iter_attn_cache, None), mem = layer_mem, mem_mask = layer_mem_mask, attn_bias = attn_bias, kv_input_residual = next(self_attn_kv_residuals_iter, None), value_residual = maybe_self_attn_value_residual, return_intermediates = True)
             elif layer_type == 'c':
                 out, inter = block(x, context = context, mask = mask, context_mask = context_mask, prev_attn = prev_cross_attn, cache = next(iter_attn_cache, None), kv_input_residual = next(cross_attn_kv_residuals_iter, None), value_residual = maybe_cross_attn_value_residual, **cross_attn_rotary_pos_emb, return_intermediates = True)
             elif layer_type == 'f':