hyper-connections 0.1.12__py3-none-any.whl → 0.1.14__py3-none-any.whl

hyper_connections/residuals.py

@@ -2,7 +2,7 @@ import torch
 from torch import nn
 from torch.nn import Module
 
-from einops import rearrange
+from einops import rearrange, pack, unpack
 
 class GRUGatedResidual(Module):
     def __init__(
@@ -13,10 +13,26 @@ class GRUGatedResidual(Module):
         self.gru = nn.GRUCell(dim, dim)
 
     def forward(self, x, residual):
+        x, ps = pack([x], '* d')
+        residual, _ = pack([residual], '* d')
 
-        gated_output = self.gru(
-            rearrange(x, 'b n d -> (b n) d'),
-            rearrange(residual, 'b n d -> (b n) d')
-        )
+        output = self.gru(x, residual)
 
-        return gated_output.reshape_as(x)
+        output, = unpack(output, ps, '* d')
+        return output
+
+class GatedResidual(Module):
+    def __init__(
+        self,
+        dim
+    ):
+        super().__init__()
+        self.to_learned_mix = nn.Linear(dim * 2, dim)
+
+    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.dist-info → hyper_connections-0.1.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.1.12
+Version: 0.1.14
 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.12.dist-info → hyper_connections-0.1.14.dist-info}/RECORD

@@ -3,8 +3,8 @@ hyper_connections/hyper_connections.py,sha256=vpipBRUGgYQ2qLBtT4Ws-myJYVdkQDkN3I
 hyper_connections/hyper_connections_channel_first.py,sha256=BojfO2dcT4jX1rlcU3kr0B6B_CjrkkS2AZU4ZXeWvh8,6769
 hyper_connections/hyper_connections_with_multi_branch_inputs.py,sha256=2JABz1slrF5_XP33L1CMNTmmixfoD464FtQpUADqneU,7806
 hyper_connections/hyper_connections_with_multi_input_streams.py,sha256=UIKceEZEBLKFL5VuP5tR1KTDeZNIJEKjFuPAwXkcp0I,11282
-hyper_connections/residuals.py,sha256=F8NMh9vHDN8Nyo0nzIm_owgR9gXPZUCNujAuduYbTiU,465
-hyper_connections-0.1.12.dist-info/METADATA,sha256=rRgT4Zhj2xUPSRPH0LKZmfJK4olGOxpxPR5ha9wEMQE,5231
-hyper_connections-0.1.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-hyper_connections-0.1.12.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
-hyper_connections-0.1.12.dist-info/RECORD,,
+hyper_connections/residuals.py,sha256=qapN4lt51qNWKa5nX7whN4xcNORxMdr3bdUwIMQPdpQ,853
+hyper_connections-0.1.14.dist-info/METADATA,sha256=7Agg3rGvMYkZEyX3n9yJPO4cNm8-9a33afV4Yc8r7WA,5231
+hyper_connections-0.1.14.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+hyper_connections-0.1.14.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
+hyper_connections-0.1.14.dist-info/RECORD,,