hyper-connections 0.0.2__py3-none-any.whl → 0.0.4__py3-none-any.whl

hyper_connections/hyper_connections.py

@@ -6,6 +6,7 @@ import torch
 from torch import nn
 from torch.nn import Module
 import torch.nn.functional as F
+from torch.utils._pytree import tree_flatten, tree_unflatten
 
 from einops import rearrange, repeat, reduce, einsum
 
@@ -113,7 +114,7 @@ class HyperConnections(Module):
 
         return output
 
-    def forward(self, residuals, **branch_kwargs):
+    def forward(self, residuals, *branch_args, **branch_kwargs):
 
         branch_input, residuals, beta = self.width_connection(residuals)
 
@@ -123,6 +124,10 @@ class HyperConnections(Module):
         if not exists(self.branch):
             return branch_input, add_residual_fn
 
-        branch_output = self.branch(branch_input, **branch_kwargs)
+        branch_output = self.branch(branch_input, *branch_args, **branch_kwargs)
 
-        return add_residual_fn(branch_output)
+        (branch_output, *rest), tree_spec = tree_flatten(branch_output)
+
+        branch_output = add_residual_fn(branch_output)
+
+        return tree_unflatten((branch_output, *rest), tree_spec)

{hyper_connections-0.0.2.dist-info → hyper_connections-0.0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.0.2
+Version: 0.0.4
 Summary: Hyper-Connections
 Project-URL: Homepage, https://pypi.org/project/hyper-connections/
 Project-URL: Repository, https://github.com/lucidrains/hyper-connections
@@ -45,7 +45,7 @@ Description-Content-Type: text/markdown
 
 ## Hyper Connections
 
-Attempt to make the multiple residual stream approach proposed by Hyper-Connections paper by Bytedance AI more accessible as a reusable library, and for following any new research in this direction.
+Attempt to make multiple residual streams, proposed in [Hyper-Connections paper](https://arxiv.org/abs/2409.19606) out of Bytedance AI lab, accessible as an easy to use library, as well as for following any new research in this direction.
 
 ## Install
 
@@ -114,7 +114,7 @@ from hyper_connections import HyperConnections
 
 expand_stream, reduce_stream = HyperConnections.get_expand_reduce_stream_functions(4)
 
-# 1. wrap your branch function
+# 1. instantiate hyper connection with correct number of streams (4 in this case)
 
 hyper_conn = HyperConnections(4, dim = 512)
 
@@ -124,11 +124,11 @@ residual = expand_stream(residual)
 
 # 3. forward your residual into hyper connection for the branch input + add residual function (learned betas)
 
-branch_input, depth_connect = hyper_conn(residual)
+branch_input, add_residual = hyper_conn(residual)
 
 branch_output = branch(branch_input)
 
-residual = depth_connect(branch_output)
+residual = add_residual(branch_output)
 
 # 4. reduce 4 streams with a summation, this has to be done after your for loop trunk
 

hyper_connections-0.0.4.dist-info/RECORD

@@ -0,0 +1,6 @@
+hyper_connections/__init__.py,sha256=xXx2Mb-dS1__UPzT-5VR1XZmyqKSSkT1DU6bAcK8jR0,73
+hyper_connections/hyper_connections.py,sha256=HbSNev1_wsJeZir5pqa-5IQYW3Ovy3UJD7A59utSF9Q,4331
+hyper_connections-0.0.4.dist-info/METADATA,sha256=uRvNtEjISVXZkI_GVfwEk_1z_w2C4uzFYSP9m8cYpWs,4676
+hyper_connections-0.0.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+hyper_connections-0.0.4.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
+hyper_connections-0.0.4.dist-info/RECORD,,

hyper_connections-0.0.2.dist-info/RECORD

@@ -1,6 +0,0 @@
-hyper_connections/__init__.py,sha256=xXx2Mb-dS1__UPzT-5VR1XZmyqKSSkT1DU6bAcK8jR0,73
-hyper_connections/hyper_connections.py,sha256=0AcrJ6O2Crnc3aiz-G5zExypcdujJHaTgmLJswdRA_c,4094
-hyper_connections-0.0.2.dist-info/METADATA,sha256=KptXVWgSC5aIh63S7cCSp4h7X7xCqxMOFRToGr6XEbo,4587
-hyper_connections-0.0.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-hyper_connections-0.0.2.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
-hyper_connections-0.0.2.dist-info/RECORD,,