x-transformers 1.23.4__py3-none-any.whl → 1.23.6__py3-none-any.whl

x_transformers/__init__.py

@@ -1,13 +1,19 @@
-import torch
-from packaging import version
-
-if version.parse(torch.__version__) >= version.parse('2.0.0'):
-    from einops._torch_specific import allow_ops_in_compiled_graph
-    allow_ops_in_compiled_graph()
-
-from x_transformers.x_transformers import XTransformer, Encoder, Decoder, CrossAttender, Attention, TransformerWrapper, ViTransformerWrapper, ContinuousTransformerWrapper
+from x_transformers.x_transformers import (
+    XTransformer,
+    Encoder,
+    Decoder,
+    CrossAttender,
+    Attention,
+    TransformerWrapper,
+    ViTransformerWrapper
+)
 
 from x_transformers.autoregressive_wrapper import AutoregressiveWrapper
 from x_transformers.nonautoregressive_wrapper import NonAutoregressiveWrapper
-from x_transformers.continuous_autoregressive_wrapper import ContinuousAutoregressiveWrapper
+
+from x_transformers.continuous_autoregressive_wrapper import (
+    ContinuousTransformerWrapper,
+    ContinuousAutoregressiveWrapper
+)
+
 from x_transformers.xl_autoregressive_wrapper import XLAutoregressiveWrapper

x_transformers/continuous_autoregressive_wrapper.py

@@ -2,11 +2,147 @@ import torch
 from torch import nn
 import torch.nn.functional as F
 
+from x_transformers.x_transformers import (
+    AttentionLayers,
+    ScaledSinusoidalEmbedding,
+    AbsolutePositionalEmbedding
+)
+
+# helper functions
+
 def exists(val):
     return val is not None
 
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if callable(d) else d
+
+# main classes
+
+class ContinuousTransformerWrapper(nn.Module):
+    def __init__(
+        self,
+        *,
+        max_seq_len,
+        attn_layers: AttentionLayers,
+        dim_in = None,
+        dim_out = None,
+        emb_dim = None,
+        max_mem_len = 0,
+        num_memory_tokens = None,
+        post_emb_norm = False,
+        emb_dropout = 0.,
+        use_abs_pos_emb = True,
+        scaled_sinu_pos_emb = False
+    ):
+        super().__init__()
+        dim = attn_layers.dim
+
+        self.max_seq_len = max_seq_len
+
+        self.max_mem_len = max_mem_len
+
+        if not (use_abs_pos_emb and not attn_layers.has_pos_emb):
+            self.pos_emb = always(0)
+        elif scaled_sinu_pos_emb:
+            self.pos_emb = ScaledSinusoidalEmbedding(dim)
+        else:
+            self.pos_emb = AbsolutePositionalEmbedding(dim, max_seq_len)
+
+        self.post_emb_norm = nn.LayerNorm(dim) if post_emb_norm else nn.Identity()
+        self.emb_dropout = nn.Dropout(emb_dropout)
+
+        # memory tokens
+
+        num_memory_tokens = default(num_memory_tokens, 0)
+        self.has_memory_tokens = num_memory_tokens > 0
+
+        if num_memory_tokens > 0:
+            self.memory_tokens = nn.Parameter(torch.randn(num_memory_tokens, dim))
+
+        # attention layers
+
+        self.attn_layers = attn_layers
+
+        # project in and out
+
+        self.project_in = nn.Linear(dim_in, dim) if exists(dim_in) else nn.Identity()
+        self.project_out = nn.Linear(dim, dim_out) if exists(dim_out) else nn.Identity()
+
+    def forward(
+        self,
+        x,
+        return_embeddings = False,
+        return_intermediates = False,
+        return_mems = False,
+        mask = None,
+        return_attn = False,
+        mems = None,
+        pos = None,
+        prepend_embeds = None,
+        **kwargs
+    ):
+        batch = x.shape[0]
+
+        x = self.project_in(x)
+        x = x + self.pos_emb(x, pos = pos)
+
+        x = self.post_emb_norm(x)
+
+        # memory tokens
+
+        if self.has_memory_tokens:
+            m = repeat(self.memory_tokens, 'm d -> b m d', b = batch)
+            x, mem_ps = pack([m, x], 'b * d')
+
+            if exists(mask):
+                num_mems = m.shape[-2]
+                mask = pad_at_dim(mask, (num_mems, 0), dim = -1, value = True)
+
+        # whether to append embeds, as in PaLI, for image embeddings
+
+        if exists(prepend_embeds):
+            _, prepend_dim = prepend_embeds.shape[1:]
+            assert prepend_dim == x.shape[-1], 'prepended embeddings need to have same dimensions as model dimensions'
+
+            x = torch.cat((prepend_embeds, x), dim = -2)
+
+        x = self.emb_dropout(x)
+
+        # attention layers
+
+        x, intermediates = self.attn_layers(x, mask = mask, mems = mems, return_hiddens = True, **kwargs)
+
+        # splice out memory tokens
+
+        if self.has_memory_tokens:
+            m, x = unpack(x, mem_ps, 'b * d')
+            intermediates.memory_tokens = m
+
+        out = self.project_out(x) if not return_embeddings else x
+
+        if return_intermediates:
+            return out, intermediates
+
+        if return_mems:
+            hiddens = intermediates.hiddens
+            new_mems = list(map(lambda t: t[..., -self.max_mem_len:, :].detach(), hiddens))
+            return out, new_mems
+
+        if return_attn:
+            attn_maps = list(map(lambda t: t.post_softmax_attn, intermediates.attn_intermediates))
+            return out, attn_maps
+
+        return out
+
 class ContinuousAutoregressiveWrapper(nn.Module):
-    def __init__(self, net, ignore_index = -100, pad_value = 0):
+    def __init__(
+        self,
+        net: ContinuousTransformerWrapper,
+        ignore_index = -100,
+        pad_value = 0
+    ):
         super().__init__()
         self.net = net
         self.max_seq_len = net.max_seq_len

x_transformers/x_transformers.py

@@ -6,7 +6,6 @@ from torch import nn, einsum, Tensor
 import torch.nn.functional as F
 
 from functools import partial, wraps
-from inspect import isfunction
 from collections import namedtuple
 from dataclasses import dataclass
 from typing import List, Callable, Optional
@@ -28,6 +27,7 @@ class LayerIntermediates:
     layer_hiddens: Optional[List[Tensor]] = None
     attn_z_loss: Optional[Tensor] = None
     mems: Optional[Tensor] = None
+    memory_tokens: Optional[Tensor] = None
 
 # helpers
 
@@ -37,7 +37,7 @@ def exists(val):
 def default(val, d):
     if exists(val):
         return val
-    return d() if isfunction(d) else d
+    return d() if callable(d) else d
 
 def cast_tuple(val, depth):
     return val if isinstance(val, tuple) else (val,) * depth
@@ -1339,7 +1339,7 @@ class ViTransformerWrapper(nn.Module):
         *,
         image_size,
         patch_size,
-        attn_layers,
+        attn_layers: Encoder,
         channels = 3,
         num_classes = None,
         post_emb_norm = False,
@@ -1347,7 +1347,6 @@ class ViTransformerWrapper(nn.Module):
         emb_dropout = 0.
     ):
         super().__init__()
-        assert isinstance(attn_layers, Encoder), 'attention layers must be an Encoder'
         assert divisible_by(image_size, patch_size), 'image dimensions must be divisible by the patch size'
         dim = attn_layers.dim
         num_patches = (image_size // patch_size) ** 2
@@ -1413,7 +1412,7 @@ class TransformerWrapper(nn.Module):
         *,
         num_tokens,
         max_seq_len,
-        attn_layers,
+        attn_layers: AttentionLayers,
         emb_dim = None,
         max_mem_len = 0,
         shift_mem_down = 0,
@@ -1430,7 +1429,6 @@ class TransformerWrapper(nn.Module):
         attn_z_loss_weight = 1e-4,
     ):
         super().__init__()
-        assert isinstance(attn_layers, AttentionLayers), 'attention layers must be one of Encoder or Decoder'
 
         dim = attn_layers.dim
         emb_dim = default(emb_dim, dim)
@@ -1576,6 +1574,8 @@ class TransformerWrapper(nn.Module):
 
             mem, x = unpack(x, mem_packed_shape, 'b * d')
 
+            intermediates.memory_tokens = mem
+
             if exists(mem_every):
                 x = rearrange(x, '(b n) m d -> b (n m) d', b = b)
 
@@ -1612,123 +1612,6 @@ class TransformerWrapper(nn.Module):
 
         return out
 
-class ContinuousTransformerWrapper(nn.Module):
-    def __init__(
-        self,
-        *,
-        max_seq_len,
-        attn_layers,
-        dim_in = None,
-        dim_out = None,
-        emb_dim = None,
-        max_mem_len = 0,
-        num_memory_tokens = None,
-        post_emb_norm = False,
-        emb_dropout = 0.,
-        use_abs_pos_emb = True,
-        scaled_sinu_pos_emb = False
-    ):
-        super().__init__()
-        assert isinstance(attn_layers, AttentionLayers), 'attention layers must be one of Encoder or Decoder'
-
-        dim = attn_layers.dim
-
-        self.max_seq_len = max_seq_len
-
-        self.max_mem_len = max_mem_len
-
-        if not (use_abs_pos_emb and not attn_layers.has_pos_emb):
-            self.pos_emb = always(0)
-        elif scaled_sinu_pos_emb:
-            self.pos_emb = ScaledSinusoidalEmbedding(dim)
-        else:
-            self.pos_emb = AbsolutePositionalEmbedding(dim, max_seq_len)
-
-        self.post_emb_norm = nn.LayerNorm(dim) if post_emb_norm else nn.Identity()
-        self.emb_dropout = nn.Dropout(emb_dropout)
-
-        # memory tokens
-
-        num_memory_tokens = default(num_memory_tokens, 0)
-        self.has_memory_tokens = num_memory_tokens > 0
-
-        if num_memory_tokens > 0:
-            self.memory_tokens = nn.Parameter(torch.randn(num_memory_tokens, dim))
-
-        # attention layers
-
-        self.attn_layers = attn_layers
-
-        # project in and out
-
-        self.project_in = nn.Linear(dim_in, dim) if exists(dim_in) else nn.Identity()
-        self.project_out = nn.Linear(dim, dim_out) if exists(dim_out) else nn.Identity()
-
-    def forward(
-        self,
-        x,
-        return_embeddings = False,
-        return_intermediates = False,
-        return_mems = False,
-        mask = None,
-        return_attn = False,
-        mems = None,
-        pos = None,
-        prepend_embeds = None,
-        **kwargs
-    ):
-        batch = x.shape[0]
-
-        x = self.project_in(x)
-        x = x + self.pos_emb(x, pos = pos)
-
-        x = self.post_emb_norm(x)
-
-        # memory tokens
-
-        if self.has_memory_tokens:
-            m = repeat(self.memory_tokens, 'm d -> b m d', b = batch)
-            x, mem_ps = pack([m, x], 'b * d')
-
-            if exists(mask):
-                num_mems = m.shape[-2]
-                mask = pad_at_dim(mask, (num_mems, 0), dim = -1, value = True)
-
-        # whether to append embeds, as in PaLI, for image embeddings
-
-        if exists(prepend_embeds):
-            _, prepend_dim = prepend_embeds.shape[1:]
-            assert prepend_dim == x.shape[-1], 'prepended embeddings need to have same dimensions as model dimensions'
-
-            x = torch.cat((prepend_embeds, x), dim = -2)
-
-        x = self.emb_dropout(x)
-
-        # attention layers
-
-        x, intermediates = self.attn_layers(x, mask = mask, mems = mems, return_hiddens = True, **kwargs)
-
-        # splice out memory tokens
-
-        if self.has_memory_tokens:
-            m, x = unpack(x, mem_ps, 'b * d')
-
-        out = self.project_out(x) if not return_embeddings else x
-
-        if return_intermediates:
-            return out, intermediates
-
-        if return_mems:
-            hiddens = intermediates.hiddens
-            new_mems = list(map(lambda t: t[..., -self.max_mem_len:, :].detach(), hiddens))
-            return out, new_mems
-
-        if return_attn:
-            attn_maps = list(map(lambda t: t.post_softmax_attn, intermediates.attn_intermediates))
-            return out, attn_maps
-
-        return out
-
 class XTransformer(nn.Module):
     def __init__(
         self,

{x_transformers-1.23.4.dist-info → x_transformers-1.23.6.dist-info}/METADATA

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

x_transformers-1.23.6.dist-info/RECORD

@@ -0,0 +1,12 @@
+x_transformers/__init__.py,sha256=LLbGkiUKu4nR4YOISh-1gPEZYs8R3I2KMwMZHMQ2YkU,538
+x_transformers/attend.py,sha256=hZcz_iijzbEqbXp2_BPEVL-1LoHXmYaHE6e6Oy-7hFE,11263
+x_transformers/autoregressive_wrapper.py,sha256=f2u0usjUfAlXwgTz87O8J8XjGTbsbrx2XEP6K2beSNI,8944
+x_transformers/continuous_autoregressive_wrapper.py,sha256=lGqE5vFaDeuLFc7b-dAQ0hx3H4dHFK_yD4-tZQZ7vqQ,5337
+x_transformers/nonautoregressive_wrapper.py,sha256=AQLE4rA_Kh8VNoe9OzpwyeWson34sRkhks4dn4seNjI,10414
+x_transformers/x_transformers.py,sha256=8zfU1iqrd6AwT-L23jkK7tpunxfjLi4HKgbZ6evBWKU,58416
+x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
+x_transformers-1.23.6.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.23.6.dist-info/METADATA,sha256=UnmCziZTo4xYZJjAdgtRyWLytboEAR0JmyK2i0hwMvc,661
+x_transformers-1.23.6.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
+x_transformers-1.23.6.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.23.6.dist-info/RECORD,,

x_transformers-1.23.4.dist-info/RECORD

@@ -1,12 +0,0 @@
-x_transformers/__init__.py,sha256=FDb654rUx8FpXRd76B8q0diH8I7q-ZjTWEtEJ4UM21Y,701
-x_transformers/attend.py,sha256=hZcz_iijzbEqbXp2_BPEVL-1LoHXmYaHE6e6Oy-7hFE,11263
-x_transformers/autoregressive_wrapper.py,sha256=f2u0usjUfAlXwgTz87O8J8XjGTbsbrx2XEP6K2beSNI,8944
-x_transformers/continuous_autoregressive_wrapper.py,sha256=pTiDqu6JRUlnQJQp_xHATYHy0lgSd6ERLqyiFO3pC-4,1575
-x_transformers/nonautoregressive_wrapper.py,sha256=AQLE4rA_Kh8VNoe9OzpwyeWson34sRkhks4dn4seNjI,10414
-x_transformers/x_transformers.py,sha256=jjuh7MLIlV4pHsJHONFlQy96QTdI5VonR9cc8FVs4J8,62009
-x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
-x_transformers-1.23.4.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.23.4.dist-info/METADATA,sha256=51dyPgX1be0-AH9PktXW_fme4GejI0YgZXbviy97j0Q,661
-x_transformers-1.23.4.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-x_transformers-1.23.4.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.23.4.dist-info/RECORD,,