titans-pytorch 0.1.6__py3-none-any.whl → 0.1.7__py3-none-any.whl

titans_pytorch/titans.py

@@ -161,14 +161,16 @@ class GatedResidualMemoryMLP(Module):
     def __init__(
         self,
         dim,
-        depth
+        depth,
+        expansion_factor = 2.
     ):
         super().__init__()
-        self.depth = depth
+        dim_hidden = int(dim * expansion_factor)
 
         self.weights = ParameterList([
             ParameterList([
-                Parameter(torch.randn(dim, dim)),
+                Parameter(torch.randn(dim, dim_hidden)),
+                Parameter(torch.randn(dim_hidden, dim)),
                 Parameter(torch.randn(dim * 2, dim)),
             ]) for _ in range(depth)
         ])
@@ -182,16 +184,17 @@ class GatedResidualMemoryMLP(Module):
         self,
         x
     ):
-        for weight, to_gates in self.weights:
+        for weight1, weight2, to_gates in self.weights:
             res = x
 
-            x = x @ weight
-            x = F.silu(x)
+            hidden = x @ weight1
+            hidden = F.silu(hidden)
+            branch_out = hidden @ weight2
 
             # gated residual
 
-            gates = cat((x, res), dim = -1) @ to_gates
-            x = res.lerp(x, gates.sigmoid())
+            gates = cat((branch_out, res), dim = -1) @ to_gates
+            x = res.lerp(branch_out, gates.sigmoid())
 
         return x @ self.final_proj
 

{titans_pytorch-0.1.6.dist-info → titans_pytorch-0.1.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.1.6
+Version: 0.1.7
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch

titans_pytorch-0.1.7.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=u0tta_KqhOdfzCEDWT9P4_jejJEK2q1XxhsEzB5MnQU,223
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=YowJzQC2p3fWgzFCe9jXrw1b3wWPKN-dcLs5sX-3r8Y,19123
+titans_pytorch/titans.py,sha256=GFnORd9WRQBBavEUWz0DV0lxnHtLsrDEghX0RgWhZaQ,20758
+titans_pytorch-0.1.7.dist-info/METADATA,sha256=ltCnRZ1MXbqxJ5-L41Z1pvfTDpLJNnk_XZ9hpfX6wrQ,4747
+titans_pytorch-0.1.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.1.7.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.1.7.dist-info/RECORD,,

titans_pytorch-0.1.6.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=u0tta_KqhOdfzCEDWT9P4_jejJEK2q1XxhsEzB5MnQU,223
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=YowJzQC2p3fWgzFCe9jXrw1b3wWPKN-dcLs5sX-3r8Y,19123
-titans_pytorch/titans.py,sha256=VMcPcKsoR3G13Um62Aa1HbdwrrV60ljPhP-yF40x90I,20555
-titans_pytorch-0.1.6.dist-info/METADATA,sha256=LJW26WfT9WB-0NfokLLHhcRpWnt76jwkXMt_FSTI3SM,4747
-titans_pytorch-0.1.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.1.6.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.1.6.dist-info/RECORD,,