hyper-connections 0.1.6__py3-none-any.whl → 0.1.7__py3-none-any.whl

hyper_connections/hyper_connections.py

@@ -73,16 +73,18 @@ class Residual(Module):
         self,
         *args,
         branch: Module | None = None,
+        residual_transform: Module | None = None,
         **kwargs
     ):
         super().__init__()
         self.branch = branch
+        self.residual_transform = default(residual_transform, nn.Identity())
 
     def width_connection(self, residuals):
         return residuals, residuals, dict()
 
     def depth_connection(self, branch_output, residuals):
-        return branch_output + residuals
+        return branch_output + self.residual_transform(residuals)
 
     def decorate_branch(self, branch: Callable):
         assert not exists(self.branch), 'branch was already wrapped on init'
@@ -128,7 +130,8 @@ class HyperConnections(Module):
         layer_index = None,
         tanh = True,
         channel_first = False,
-        dropout = 0.
+        dropout = 0.,
+        residual_transform: Module | None = None, # to support resnet blocks where dimension in not equal to dimension out - usually a residual conv
     ):
         """
         Appendix J, Algorithm2 in - https://arxiv.org/abs/2409.19606
@@ -168,7 +171,14 @@ class HyperConnections(Module):
 
         self.channel_first = channel_first
 
+        # maybe residual transform
+
+        self.residual_transform = default(residual_transform, nn.Identity())
+
     def width_connection(self, residuals):
+
+        maybe_transformed_residuals = self.residual_transform(residuals)
+
         # width connection
 
         if self.channel_first:
@@ -197,7 +207,7 @@ class HyperConnections(Module):
         if self.channel_first:
             branch_input = rearrange(branch_input, 'b ... d -> b d ...')
 
-        return branch_input, residuals, dict(beta = beta)
+        return branch_input, maybe_transformed_residuals, dict(beta = beta)
 
     def depth_connection(self, branch_output, residuals, *, beta):
         # 'depth' connection
@@ -205,13 +215,15 @@ class HyperConnections(Module):
         if self.channel_first:
             branch_output = rearrange(branch_output, 'b d ... -> b ... d')
 
-        residuals = einsum(branch_output, beta, 'b ... d, b ... s -> b ... s d') + residuals
-        output = rearrange(residuals, 'b ... s d -> (b s) ... d')
+        output = einsum(branch_output, beta, 'b ... d, b ... s -> b ... s d')
+        output = rearrange(output, 'b ... s d -> (b s) ... d')
 
         if self.channel_first:
             output = rearrange(output, 'b ... d -> b d ...')
 
-        return self.dropout(output)
+        residuals = residuals + output
+
+        return self.dropout(residuals)
 
     def decorate_branch(self, branch: Callable):
         assert not exists(self.branch), 'branch was already wrapped on init'

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

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

@@ -1,8 +1,8 @@
 hyper_connections/__init__.py,sha256=d2zNTka0Gp9vINu4U-RhgTJFBhsVrs1fne_15Zl0oOs,224
-hyper_connections/hyper_connections.py,sha256=KBf5xiblM53jbKp_IwAgWAKQetLYWKmK8-6qNn2fHoU,10007
+hyper_connections/hyper_connections.py,sha256=L2e4DduzPGdH30NhfHuiSiVZTwXRgeZW2MDAZ0Z-TKk,10541
 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.6.dist-info/METADATA,sha256=3KIWntaZvm8k5BCCs_L1lOmmpB03jgEyVy8y9kYAZxI,5230
-hyper_connections-0.1.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-hyper_connections-0.1.6.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
-hyper_connections-0.1.6.dist-info/RECORD,,
+hyper_connections-0.1.7.dist-info/METADATA,sha256=YThD719ySS2H6ABQnrNHKmoWI9vaGh5d1H9mbMKehV0,5230
+hyper_connections-0.1.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+hyper_connections-0.1.7.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
+hyper_connections-0.1.7.dist-info/RECORD,,