hyper-connections 0.1.12__tar.gz → 0.1.15__tar.gz

{hyper_connections-0.1.12 → hyper_connections-0.1.15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.1.12
+Version: 0.1.15
 Summary: Hyper-Connections
 Project-URL: Homepage, https://pypi.org/project/hyper-connections/
 Project-URL: Repository, https://github.com/lucidrains/hyper-connections

hyper_connections-0.1.15/hyper_connections/residuals.py

@@ -0,0 +1,40 @@
+import torch
+from torch import nn
+from torch.nn import Module
+
+from einops import rearrange, pack, unpack
+
+class GRUGatedResidual(Module):
+    def __init__(
+        self,
+        dim
+    ):
+        super().__init__()
+        self.gru = nn.GRUCell(dim, dim)
+
+    def forward(self, x, residual):
+        x, packed_shape = pack([x], '* d')
+        residual, _ = pack([residual], '* d')
+
+        output = self.gru(x, residual)
+
+        output, = unpack(output, packed_shape, '* d')
+        return output
+
+class GatedResidual(Module):
+    def __init__(
+        self,
+        dim,
+        fine_gate = False
+    ):
+        super().__init__()
+
+        self.to_learned_mix = nn.Linear(dim * 2, dim if fine_gate else 1)
+
+    def forward(self, x, residual):
+        x_and_residual, _ = pack([x, residual], 'b n *')
+
+        mix = self.to_learned_mix(x_and_residual)
+
+        out = x.lerp(residual, mix.sigmoid())
+        return out

{hyper_connections-0.1.12 → hyper_connections-0.1.15}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hyper-connections"
-version = "0.1.12"
+version = "0.1.15"
 description = "Hyper-Connections"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

hyper_connections-0.1.12/hyper_connections/residuals.py

@@ -1,22 +0,0 @@
-import torch
-from torch import nn
-from torch.nn import Module
-
-from einops import rearrange
-
-class GRUGatedResidual(Module):
-    def __init__(
-        self,
-        dim
-    ):
-        super().__init__()
-        self.gru = nn.GRUCell(dim, dim)
-
-    def forward(self, x, residual):
-
-        gated_output = self.gru(
-            rearrange(x, 'b n d -> (b n) d'),
-            rearrange(residual, 'b n d -> (b n) d')
-        )
-
-        return gated_output.reshape_as(x)