hyper-connections 0.0.2__py3-none-any.whl → 0.0.3__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
 
@@ -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.dist-info → hyper_connections-0.0.3.dist-info}/METADATA

@@ -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.3.dist-info/RECORD

@@ -0,0 +1,6 @@
+hyper_connections/__init__.py,sha256=xXx2Mb-dS1__UPzT-5VR1XZmyqKSSkT1DU6bAcK8jR0,73
+hyper_connections/hyper_connections.py,sha256=91QtTtnpffmErIZvrnTtosSf4JgBqcyGvxftmka-EOw,4303
+hyper_connections-0.0.3.dist-info/METADATA,sha256=8XKTmC6Ys10uOyotPtvL17v4uZyepkWoeMRVM4B_TSQ,4676
+hyper_connections-0.0.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+hyper_connections-0.0.3.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
+hyper_connections-0.0.3.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,,