hyper-connections 0.0.10__tar.gz → 0.0.12__tar.gz

{hyper_connections-0.0.10 → hyper_connections-0.0.12}/PKG-INFO

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

{hyper_connections-0.0.10 → hyper_connections-0.0.12}/hyper_connections/hyper_connections.py

@@ -43,23 +43,37 @@ class Residual(Module):
     def depth_connection(self, branch_output, residuals):
         return branch_output + residuals
 
+    def decorate_branch(self, branch: Callable):
+        assert not exists(self.branch), 'branch was already wrapped on init'
+
+        def forward_and_add_residual(residual, *args, **kwargs):
+            branch_input, add_residual = self.forward(residual)
+
+            branch_output = branch(branch_input, *args, **kwargs)
+
+            residual = add_residual(branch_output)
+
+            return residual
+
+        return forward_and_add_residual
+
     def forward(self, residuals, *branch_args, **branch_kwargs):
 
         branch_input, residuals, residual_kwargs = self.width_connection(residuals)
 
         def add_residual_fn(branch_out):
-            return self.depth_connection(branch_out, residuals, **residual_kwargs)
+            (branch_out, *rest), tree_spec = tree_flatten(branch_out)
+
+            branch_out = self.depth_connection(branch_out, residuals, **residual_kwargs)
+
+            return tree_unflatten((branch_out, *rest), tree_spec)
 
         if not exists(self.branch):
             return branch_input, add_residual_fn
 
         branch_output = self.branch(branch_input, *branch_args, **branch_kwargs)
 
-        (branch_output, *rest), tree_spec = tree_flatten(branch_output)
-
-        branch_output = add_residual_fn(branch_output)
-
-        return tree_unflatten((branch_output, *rest), tree_spec)
+        return add_residual_fn(branch_output)
 
 # hyper connection residual streams
 
@@ -174,7 +188,7 @@ class HyperConnections(Module):
         def forward_and_add_residual(residual, *args, **kwargs):
             branch_input, add_residual = self.forward(residual)
 
-            branch_output = branch(branch_input)
+            branch_output = branch(branch_input, *args, **kwargs)
 
             residual = add_residual(branch_output)
 
@@ -187,18 +201,18 @@ class HyperConnections(Module):
         branch_input, residuals, residual_kwargs = self.width_connection(residuals)
 
         def add_residual_fn(branch_out):
-            return self.depth_connection(branch_out, residuals, **residual_kwargs)
+            (branch_out, *rest), tree_spec = tree_flatten(branch_out)
+
+            branch_out = self.depth_connection(branch_out, residuals, **residual_kwargs)
+
+            return tree_unflatten((branch_out, *rest), tree_spec)
 
         if not exists(self.branch):
             return branch_input, add_residual_fn
 
         branch_output = self.branch(branch_input, *branch_args, **branch_kwargs)
 
-        (branch_output, *rest), tree_spec = tree_flatten(branch_output)
-
-        branch_output = add_residual_fn(branch_output)
-
-        return tree_unflatten((branch_output, *rest), tree_spec)
+        return add_residual_fn(branch_output)
 
 # stream embed
 

{hyper_connections-0.0.10 → hyper_connections-0.0.12}/pyproject.toml

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