x-transformers 1.41.5__py3-none-any.whl → 1.42.0__py3-none-any.whl

x_transformers/__init__.py

@@ -32,3 +32,7 @@ from x_transformers.xl_autoregressive_wrapper import XLAutoregressiveWrapper
 from x_transformers.dpo import (
     DPO
 )
+
+from x_transformers.neo_mlp import (
+    NeoMLP
+)

x_transformers/neo_mlp.py

@@ -0,0 +1,126 @@
+from collections import namedtuple
+
+import torch
+from torch import nn, tensor, pi, is_tensor
+import torch.nn.functional as F
+from torch.nn import Module, ModuleList
+
+from einops import rearrange, repeat, einsum, pack, unpack
+
+from x_transformers.x_transformers import (
+    Encoder
+)
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+# random fourier
+
+class RandomFourierEmbed(Module):
+
+    def __init__(self, dim):
+        super().__init__()
+        self.proj = nn.Linear(1, dim)
+        self.proj.requires_grad_(False)
+
+    def forward(
+        self,
+        times,
+    ):
+
+        times = rearrange(times, '... -> ... 1')
+        rand_proj = self.proj(times)
+        return torch.cos(2 * pi * rand_proj)
+
+# class
+
+class NeoMLP(Module):
+    """ https://openreview.net/forum?id=A8Vuf2e8y6 """
+
+    def __init__(
+        self,
+        *,
+        dim_in,
+        dim_hidden,
+        dim_out,
+        dim_model,
+        depth,
+        encoder_kwargs: dict = dict(
+            attn_dim_head = 16,
+            heads = 4
+        )
+    ):
+        super().__init__()
+
+        # input and output embeddings
+
+        self.input_embed = nn.Parameter(torch.zeros(dim_in, dim_model))
+        self.hidden_embed = nn.Parameter(torch.zeros(dim_hidden, dim_model))
+        self.output_embed = nn.Parameter(torch.zeros(dim_out, dim_model))
+
+        nn.init.normal_(self.input_embed, std = 0.02)
+        nn.init.normal_(self.hidden_embed, std = 0.02)
+        nn.init.normal_(self.output_embed, std = 0.02)
+
+        # they use random fourier for continuous features
+
+        self.random_fourier = nn.Sequential(
+            RandomFourierEmbed(dim_model),
+            nn.Linear(dim_model, dim_model)
+        )
+
+        # hidden dimensions of mlp replaced with nodes with message passing
+        # which comes back to self attention as a fully connected graph.
+
+        self.transformer = Encoder(
+            dim = dim_model,
+            depth = depth,
+            **encoder_kwargs
+        )
+
+        # output
+
+        self.to_output_weights = nn.Parameter(torch.randn(dim_out, dim_model))
+        self.to_output_bias = nn.Parameter(torch.zeros(dim_out))
+
+    def forward(
+        self,
+        x,
+        return_embeds = False
+    ):
+        batch = x.shape[0]
+
+        fouriered_input = self.random_fourier(x)
+
+        # add fouriered input to the input embedding
+
+        input_embed = fouriered_input + self.input_embed
+
+        hidden_embed, output_embed = tuple(repeat(t, '... -> b ...', b = batch) for t in (self.hidden_embed, self.output_embed))
+
+        # pack all the inputs into one string of tokens for self attention
+
+        embed, packed_shape = pack([input_embed, hidden_embed, output_embed], 'b * d')
+
+        # attention is all you need
+
+        embed = self.transformer(embed)
+
+        # unpack
+
+        input_embed, hidden_embed, output_embed = unpack(embed, packed_shape, 'b * d')
+
+        # project for output
+
+        output = einsum(output_embed, self.to_output_weights, 'b n d, n d -> b n')
+        output = output + self.to_output_bias
+
+        if not return_embeds:
+            return output
+
+        return output, (input_embed, hidden_embed, output_embed)

{x_transformers-1.41.5.dist-info → x_transformers-1.42.0.dist-info}/METADATA

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

{x_transformers-1.41.5.dist-info → x_transformers-1.42.0.dist-info}/RECORD

@@ -1,15 +1,16 @@
-x_transformers/__init__.py,sha256=-MkQrSc37cTVDX7AOykxunYnqVtFlQ7lb0Cse5dsGWU,793
+x_transformers/__init__.py,sha256=l0dom8ZYkRzFvnDdgzDboXqrI1tKav3beVE7TN2nHko,844
 x_transformers/attend.py,sha256=SdWlV8Vp5DtpsOzAd0LRhm4VGrJf0lJCGiV2_j_CtoA,17284
 x_transformers/autoregressive_wrapper.py,sha256=DOJJCMMDOqDYKWy_IaG5IyKsXD3AW6amzfUgdAADOLY,10500
 x_transformers/continuous.py,sha256=cIVEdhfei258__ziV7kQBrJMxCel54bExBTDrO9rfCI,6450
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
+x_transformers/neo_mlp.py,sha256=nNQCbNM_uxBS_oc4J28BXDVpsvzyFJIbgm9xgEdjL0c,3221
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
 x_transformers/x_transformers.py,sha256=UhIbFPXjdQsbFBDHVGmV81LGHTD5qwbusDc5kl3F2A4,91987
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-1.41.5.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.41.5.dist-info/METADATA,sha256=4op6TctcnQVjp6pWaYugV7rxG-5e-pn6wL_qS95d98E,689
-x_transformers-1.41.5.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-x_transformers-1.41.5.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.41.5.dist-info/RECORD,,
+x_transformers-1.42.0.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.42.0.dist-info/METADATA,sha256=cwUbK81bkaChT3ZTyqRr7biJx87s394uJAQPn0df_38,689
+x_transformers-1.42.0.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+x_transformers-1.42.0.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.42.0.dist-info/RECORD,,