hyper-connections 0.0.3__py3-none-any.whl → 0.0.5__py3-none-any.whl

hyper_connections/hyper_connections.py

@@ -69,6 +69,14 @@ class HyperConnections(Module):
 
         return expand_fn, reduce_fn
 
+    @classmethod
+    def get_init_and_expand_reduce_stream_functions(cls, num_streams):
+
+        init_hyper_conn_fn = partial(cls, num_streams)
+        expand_reduce_fns = cls.get_expand_reduce_stream_functions(num_streams)
+
+        return (init_hyper_conn_fn, *expand_reduce_fns)
+
     def width_connection(self, residuals):
         # width connection
 
@@ -114,7 +122,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)
 
@@ -124,7 +132,7 @@ 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)
 
         (branch_output, *rest), tree_spec = tree_flatten(branch_output)
 

{hyper_connections-0.0.3.dist-info → hyper_connections-0.0.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.0.3
+Version: 0.0.5
 Summary: Hyper-Connections
 Project-URL: Homepage, https://pypi.org/project/hyper-connections/
 Project-URL: Repository, https://github.com/lucidrains/hyper-connections
@@ -73,11 +73,11 @@ residual = branch(residual) + residual
 
 from hyper_connections import HyperConnections
 
-expand_stream, reduce_stream = HyperConnections.get_expand_reduce_stream_functions(4)
+init_hyper_conn, expand_stream, reduce_stream = HyperConnections.get_init_and_expand_reduce_stream_functions(4)
 
 # 1. wrap your branch function
 
-hyper_conn_branch = HyperConnections(4, dim = 512, branch = branch)
+hyper_conn_branch = init_hyper_conn(dim = 512, branch = branch)
 
 # 2. expand to 4 streams, this must be done before your trunk, typically a for-loop with many branch functions
 
@@ -112,11 +112,11 @@ residual = branch(residual) + residual
 
 from hyper_connections import HyperConnections
 
-expand_stream, reduce_stream = HyperConnections.get_expand_reduce_stream_functions(4)
+init_hyper_conn, expand_stream, reduce_stream = HyperConnections.get_init_and_expand_reduce_stream_functions(4)
 
-# 1. instantiate hyper connection with correct number of streams (4 in this case)
+# 1. instantiate hyper connection with correct number of streams (4 in this case) - or use the init function above
 
-hyper_conn = HyperConnections(4, dim = 512)
+hyper_conn = init_hyper_conn(dim = 512)
 
 # 2. expand to 4 streams
 

hyper_connections-0.0.5.dist-info/RECORD

@@ -0,0 +1,6 @@
+hyper_connections/__init__.py,sha256=xXx2Mb-dS1__UPzT-5VR1XZmyqKSSkT1DU6bAcK8jR0,73
+hyper_connections/hyper_connections.py,sha256=2lZcPuW4hEKet3r8caN-sN-PzRaBNL1q-V3_uA1lVaM,4613
+hyper_connections-0.0.5.dist-info/METADATA,sha256=bko4lEEBiulROYd0aOC_8rIk0hPinBg5TBQWf9DZe9M,4753
+hyper_connections-0.0.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+hyper_connections-0.0.5.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
+hyper_connections-0.0.5.dist-info/RECORD,,

hyper_connections-0.0.3.dist-info/RECORD

@@ -1,6 +0,0 @@
-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,,