x-transformers 1.42.7__tar.gz → 1.42.9__tar.gz

{x_transformers-1.42.7/x_transformers.egg-info → x_transformers-1.42.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.42.7
+Version: 1.42.9
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang

{x_transformers-1.42.7 → x_transformers-1.42.9}/setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'x-transformers',
   packages = find_packages(exclude=['examples']),
-  version = '1.42.7',
+  version = '1.42.9',
   license='MIT',
   description = 'X-Transformers - Pytorch',
   author = 'Phil Wang',

{x_transformers-1.42.7 → x_transformers-1.42.9}/tests/test_x_transformers.py

@@ -1,12 +1,15 @@
 import pytest
+
 import torch
+from torch import nn
+from torch.nn import Module
 
 from x_transformers.x_transformers import (
     XTransformer,
     TransformerWrapper,
     Encoder,
     Decoder,
-    AutoregressiveWrapper,
+    LinearNoBias,
 )
 
 from x_transformers.neo_mlp import (
@@ -378,6 +381,7 @@ def test_neo_mlp():
     assert out.shape == (3, 7)
 
 def test_custom_alibi():
+
     model = TransformerWrapper(
         num_tokens = 20_000,
         max_seq_len = 1024,
@@ -394,3 +398,113 @@ def test_custom_alibi():
     pos = torch.tensor([[0, 1, 2, 4], [1, 3, 5, 7]])
 
     logits = model(x, pos = pos)
+
+def test_custom_alibi_across_heads():
+
+    model = Decoder(
+        dim = 512,
+        depth = 2,
+        heads = 2,
+        alibi_pos_bias = True,
+        rel_pos_kwargs = dict(
+            slopes = [1, 1]
+        ),
+    )
+
+    x = torch.randn(2, 4, 512)
+
+    pos = torch.tensor([
+        [[0, 1, 2, 4], [1, 3, 5, 7]],
+        [[2, 3, 4, 5], [6, 8, 9, 10]]
+    ])
+
+    embed = model(x, pos = pos)
+
+@pytest.mark.parametrize('embedder_type', ('embedding', 'none', 'custom'))
+def test_embedder(embedder_type):
+    num_tokens = 20000
+    dim = 128
+    token_emb_kwargs = {}
+
+    if embedder_type == 'embedding':
+        embedder = nn.Embedding(num_tokens, dim)
+    elif embedder_type == 'none':
+        embedder = None
+    else:
+        class CustomEmbedder(Module):
+            """
+            Made up embedder that sums two embeddings. Just to check if we can pass additional input to the embedder's
+            forward pass without breaking the model.
+            """
+            def __init__(self, num_tokens, dim):
+                super().__init__()
+                self.embed_x = nn.Embedding(num_tokens, dim)
+                self.embed_y = nn.Embedding(num_tokens, dim)
+
+            def forward(self, x, y):
+                return self.embed_x(x) + self.embed_y(y)
+
+            def init_(self):
+                pass
+
+        embedder = CustomEmbedder(num_tokens, dim)
+        token_emb_kwargs['y'] = torch.randint(0, num_tokens, (2, 1024))
+
+    model = TransformerWrapper(
+        num_tokens = num_tokens,
+        max_seq_len = 1024,
+        attn_layers = Decoder(
+            dim = dim,
+            depth = 6,
+            heads = 8,
+        ),
+        token_emb = embedder,
+    )
+
+    x = torch.randint(0, 20000, (2, 1024))
+
+    output = model(x, token_emb_kwargs=token_emb_kwargs)
+    assert output.shape == (2, 1024, 20000)
+
+
+@pytest.mark.parametrize("to_logits", ('linear', 'none', 'pointer'))
+def test_to_logits(to_logits):
+    num_tokens = 20000
+    dim = 128
+
+    to_logits_kwargs = {}
+
+    if to_logits == 'linear':
+        logit_mapper = LinearNoBias(dim, num_tokens)
+    elif to_logits == 'none':
+        logit_mapper = None
+    else:
+        class PointerNetworkLogits(Module):
+            def __init__(self, dim):
+                super().__init__()
+                self.proj_to_pointers = nn.Linear(dim, dim)
+
+            def forward(self, model_embeddings, input_embeddings):
+                pointers = self.proj_to_pointers(model_embeddings)
+                logits = torch.matmul(pointers, input_embeddings.permute(0, 2, 1))
+                return logits
+
+        logit_mapper = PointerNetworkLogits(dim)
+        to_logits_kwargs['input_embeddings'] = torch.randn(2, 20000, dim)
+
+    model = TransformerWrapper(
+        num_tokens = num_tokens,
+        max_seq_len = 1024,
+        attn_layers = Decoder(
+            dim = dim,
+            depth = 6,
+            heads = 8,
+        ),
+        to_logits = logit_mapper,
+    )
+
+    x = torch.randint(0, num_tokens, (2, 1024))
+
+    output = model(x, to_logits_kwargs=to_logits_kwargs)
+
+    assert output.shape == (2, 1024, 20000)

{x_transformers-1.42.7 → x_transformers-1.42.9}/x_transformers/x_transformers.py

@@ -238,6 +238,13 @@ class TokenEmbedding(Module):
         token_emb = self.emb(x.long())
         return l2norm(token_emb) if self.l2norm_embed else token_emb
 
+    def init_(self):
+        if self.l2norm_embed:
+            nn.init.normal_(self.emb.weight, std=1e-5)
+            return
+        nn.init.kaiming_normal_(self.emb.weight)
+
+
 # positional embeddings
 
 class AbsolutePositionalEmbedding(Module):
@@ -445,13 +452,20 @@ class DynamicPositionBias(Module):
         return bias
 
 class AlibiPositionalBias(Module):
-    def __init__(self, heads, total_heads = None, **kwargs):
+    def __init__(
+        self,
+        heads,
+        total_heads = None,
+        slopes: list[int] | None = None,
+        **kwargs
+    ):
         super().__init__()
         self.heads = heads
         self.total_heads = default(total_heads, heads)
 
-        slopes = Tensor(self._get_slopes(heads))
+        slopes = Tensor(default(slopes, self._get_slopes(heads)))
         slopes = rearrange(slopes, 'h -> h 1 1')
+
         self.register_buffer('slopes', slopes, persistent = False)
         self.register_buffer('bias', None, persistent = False)
     
@@ -480,7 +494,10 @@ class AlibiPositionalBias(Module):
         h, device = self.total_heads, self.device
 
         pos_j = default(pos_j, pos_i)
-        bias = -einx.subtract('... j, ... i -> ... 1 i j', pos_j, pos_i).abs()
+        bias = -einx.subtract('... j, ... i -> ... i j', pos_j, pos_i).abs()
+
+        if bias.ndim == 3:
+            bias = rearrange(bias, 'b i j -> b 1 i j')
 
         bias = bias * self.slopes
         num_heads_unalibied = h - bias.shape[-3]
@@ -1524,8 +1541,9 @@ class AttentionLayers(Module):
         use_layerscale = False,
         layerscale_init_value = 0.,
         unet_skips = False,
-        reinject_input = False, # seen first in DEQ paper https://arxiv.org/abs/1909.01377, but later used in a number of papers trying to achieve depthwise generalization https://arxiv.org/abs/2410.03020v1
-        add_value_residual = False, # resformer from Zhou et al - https://arxiv.org/abs/2410.17897v1
+        reinject_input = False,         # seen first in DEQ paper https://arxiv.org/abs/1909.01377, but later used in a number of papers trying to achieve depthwise generalization https://arxiv.org/abs/2410.03020v1
+        add_value_residual = False,     # resformer from Zhou et al - https://arxiv.org/abs/2410.17897v1
+        rel_pos_kwargs: dict = dict(),
         **kwargs
     ):
         super().__init__()
@@ -1566,14 +1584,14 @@ class AttentionLayers(Module):
 
         if rel_pos_bias:
             assert not flash_attn, 'flash attention not compatible with t5 relative positional bias'
-            self.rel_pos = RelativePositionBias(scale = dim_head ** 0.5, causal = causal, heads = heads, num_buckets = rel_pos_num_buckets, max_distance = rel_pos_max_distance)
+            self.rel_pos = RelativePositionBias(scale = dim_head ** 0.5, causal = causal, heads = heads, num_buckets = rel_pos_num_buckets, max_distance = rel_pos_max_distance, **rel_pos_kwargs)
         elif dynamic_pos_bias:
             assert not flash_attn, 'flash attention not compatible with dynamic positional bias'
-            self.rel_pos = DynamicPositionBias(dim = dim // 4, heads = heads, log_distance = dynamic_pos_bias_log_distance, depth = dynamic_pos_bias_mlp_depth, norm = dynamic_pos_bias_norm)
+            self.rel_pos = DynamicPositionBias(dim = dim // 4, heads = heads, log_distance = dynamic_pos_bias_log_distance, depth = dynamic_pos_bias_mlp_depth, norm = dynamic_pos_bias_norm, **rel_pos_kwargs)
         elif alibi_pos_bias:
             alibi_num_heads = default(alibi_num_heads, heads)
             assert alibi_num_heads <= heads, 'number of ALiBi heads must be less than the total number of heads'
-            self.rel_pos = AlibiPositionalBias(heads = alibi_num_heads, total_heads = heads)
+            self.rel_pos = AlibiPositionalBias(heads = alibi_num_heads, total_heads = heads, **rel_pos_kwargs)
 
         assert at_most_one_of(sandwich_norm, resi_dual), 'either sandwich norm or resiDual is selected, but not both'
         assert not (not pre_norm and sandwich_norm), 'sandwich norm cannot be used when not using prenorm'
@@ -2261,7 +2279,8 @@ class TransformerWrapper(Module):
         token_emb: TokenEmbedding | None = None,
         mixture_of_softmax = False,
         mixture_of_softmax_k = 4,
-        sigsoftmax_logits = False
+        sigsoftmax_logits = False,
+        to_logits: Module | None = None,
     ):
         super().__init__()
 
@@ -2363,11 +2382,12 @@ class TransformerWrapper(Module):
         if return_only_embed:
             self.to_logits = None
         elif tie_embedding:
+            assert isinstance(token_emb, TokenEmbedding), 'can only tie embedding if using `TokenEmbedding`'
             self.to_logits = lambda t: t @ self.token_emb.emb.weight.t()
         elif num_output_heads > 1:
             self.to_logits = ModuleList([LinearNoBias(dim, logits_dim) for _ in range(num_output_heads)])
         else:
-            self.to_logits = LinearNoBias(dim, logits_dim)
+            self.to_logits = LinearNoBias(dim, logits_dim) if not exists(to_logits) else to_logits
 
         # memory tokens (like [cls]) from Memory Transformers paper
 
@@ -2388,13 +2408,12 @@ class TransformerWrapper(Module):
         self.can_cache_kv_outside_max_seq_len = no_abs_pos_emb
 
     def init_(self):
+        if hasattr(self.token_emb, 'init_'):
+            self.token_emb.init_()
+
         if self.l2norm_embed:
-            nn.init.normal_(self.token_emb.emb.weight, std = 1e-5)
             if not isinstance(self.pos_emb, always):
                 nn.init.normal_(self.pos_emb.emb.weight, std = 1e-5)
-            return
-
-        nn.init.kaiming_normal_(self.token_emb.emb.weight)
 
     def forward(
         self,
@@ -2417,7 +2436,9 @@ class TransformerWrapper(Module):
         attn_z_loss_weight = 1e-4,
         seq_start_pos = None,
         cache: LayerIntermediates | None = None,
-        **kwargs
+        token_emb_kwargs = dict(),
+        to_logits_kwargs = dict(),
+        **kwargs,
     ):
         b, n, device, num_mems, has_memory_tokens, emb_frac_gradient, orig_mask = x.shape[0], x.shape[1], x.device, self.num_memory_tokens, self.num_memory_tokens > 0, self.emb_frac_gradient, mask
 
@@ -2428,7 +2449,7 @@ class TransformerWrapper(Module):
 
         external_pos_emb = exists(pos) and pos.dtype != torch.long
         pos_emb = self.pos_emb(x, pos = pos, seq_start_pos = seq_start_pos) if not external_pos_emb else pos
-        x = self.token_emb(x) + pos_emb
+        x = self.token_emb(x, **token_emb_kwargs) + pos_emb
 
         # add additional embeddings
 
@@ -2583,9 +2604,9 @@ class TransformerWrapper(Module):
 
         if not return_embeddings:
             if self.has_multiple_heads:
-                logits = tuple(fn(x) for fn in self.to_logits)
+                logits = tuple(fn(x, **to_logits_kwargs) for fn in self.to_logits)
             else:
-                logits = self.to_logits(x)
+                logits = self.to_logits(x, **to_logits_kwargs)
 
         # maybe sig softmax
 

{x_transformers-1.42.7 → x_transformers-1.42.9/x_transformers.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.42.7
+Version: 1.42.9
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang