hyper-connections 0.0.9__tar.gz → 0.0.11__tar.gz

{hyper_connections-0.0.9 → hyper_connections-0.0.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.0.9
+Version: 0.0.11
 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.9 → hyper_connections-0.0.11}/hyper_connections/hyper_connections.py

@@ -1,4 +1,6 @@
 from __future__ import annotations
+from typing import Callable
+
 from functools import partial
 from random import randrange
 
@@ -172,7 +174,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)
 
@@ -185,18 +187,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.9 → hyper_connections-0.0.11}/pyproject.toml

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