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

hyper_connections/hyper_connections.py

@@ -33,6 +33,9 @@ def default(v, d):
 def identity(t):
     return t
 
+def add(x, y):
+    return x + y
+
 # main functions
 
 def get_expand_reduce_stream_functions(num_streams, add_stream_embed = False, dim = None, disable = False):
@@ -143,7 +146,8 @@ class HyperConnections(Module):
         dropout = 0.,
         residual_transform: Module | None = None, # to support resnet blocks where dimension in not equal to dimension out - usually a residual conv
         add_branch_out_to_residual = True, # will disable depth connections (weighted residual sum with beta) if set False
-        num_input_views = 1 # allow for the branch module to receive multiple input views, dimension placed on the very left (before batch)
+        num_input_views = 1, # allow for the branch module to receive multiple input views, dimension placed on the very left (before batch)
+        depth_residual_fn = add
     ):
         """
         Appendix J, Algorithm2 in - https://arxiv.org/abs/2409.19606
@@ -199,6 +203,12 @@ class HyperConnections(Module):
 
         self.residual_transform = default(residual_transform, nn.Identity())
 
+        # maybe custom depth connection residual function
+        # this is to prepare for gating the addition of the branch outputs to the residual streams
+        # needed for memory lanes a la RMT / LMM
+
+        self.depth_residual_fn = depth_residual_fn
+
     def width_connection(self, residuals):
 
         maybe_transformed_residuals = self.residual_transform(residuals)
@@ -254,7 +264,7 @@ class HyperConnections(Module):
         if self.channel_first:
             output = rearrange(output, 'b ... d -> b d ...')
 
-        residuals = residuals + output
+        residuals = self.depth_residual_fn(output, residuals)
 
         return self.dropout(residuals)
 

hyper_connections/residuals.py

@@ -0,0 +1,38 @@
+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, ps = pack([x], '* d')
+        residual, _ = pack([residual], '* d')
+
+        output = self.gru(x, residual)
+
+        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.11.dist-info → hyper_connections-0.1.14.dist-info}/METADATA

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

@@ -1,9 +1,10 @@
 hyper_connections/__init__.py,sha256=d2zNTka0Gp9vINu4U-RhgTJFBhsVrs1fne_15Zl0oOs,224
-hyper_connections/hyper_connections.py,sha256=mUImPtaTE8Paygs-6vq7l_mlph1CkU__jRcE4TFim_Y,12137
+hyper_connections/hyper_connections.py,sha256=vpipBRUGgYQ2qLBtT4Ws-myJYVdkQDkN3IkpTMkxRxc,12485
 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-0.1.11.dist-info/METADATA,sha256=Ck3udilJMrT1ABRkqkNhfEdkjQXAKFBtsUAAalEu3No,5231
-hyper_connections-0.1.11.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-hyper_connections-0.1.11.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
-hyper_connections-0.1.11.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,,