hyper-connections 0.0.2__tar.gz → 0.0.3__tar.gz

{hyper_connections-0.0.2 → hyper_connections-0.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.0.2
+Version: 0.0.3
 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.2 → hyper_connections-0.0.3}/README.md

@@ -2,7 +2,7 @@
 
 ## 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
 
@@ -71,7 +71,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)
 
@@ -81,11 +81,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.2 → hyper_connections-0.0.3}/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
 
@@ -125,4 +126,8 @@ class HyperConnections(Module):
 
         branch_output = self.branch(branch_input, **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 → hyper_connections-0.0.3}/pyproject.toml

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