hyper-connections 0.0.14__tar.gz → 0.0.16__tar.gz

{hyper_connections-0.0.14 → hyper_connections-0.0.16}/PKG-INFO

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

@@ -121,7 +121,11 @@ class HyperConnections(Module):
         self.channel_first = channel_first
 
     @classmethod
-    def get_expand_reduce_stream_functions(cls, num_streams):
+    def get_expand_reduce_stream_functions(cls, num_streams, disable = False):
+
+        if disable:
+            return (identity, identity)
+
         expand_fn = partial(repeat, pattern = 'b ... -> (b s) ...', s = num_streams)
         reduce_fn = partial(reduce, pattern = '(b s) ... -> b ...', reduction = 'sum', s = num_streams)
 
@@ -133,7 +137,7 @@ class HyperConnections(Module):
         hyper_conn_klass = cls if not disable else Residual
 
         init_hyper_conn_fn = partial(hyper_conn_klass, num_streams)
-        expand_reduce_fns = cls.get_expand_reduce_stream_functions(num_streams) if not disable else (identity, identity)
+        expand_reduce_fns = cls.get_expand_reduce_stream_functions(num_streams, disable = disable)
 
         return (init_hyper_conn_fn, *expand_reduce_fns)
 

{hyper_connections-0.0.14 → hyper_connections-0.0.16}/hyper_connections/hyper_connections_with_multi_branch_inputs.py

@@ -57,12 +57,19 @@ class HyperConnections(Module):
         self.num_residual_streams = num_residual_streams
         self.num_branch_inputs = num_branch_inputs
 
-        init_residual_index = default(layer_index, randrange(num_residual_streams)) % num_residual_streams # just choose one random residual stream if layer index not given
-
         self.static_beta = nn.Parameter(torch.ones(num_residual_streams, num_branch_inputs))
 
-        init_alpha0 = torch.zeros((num_residual_streams, num_branch_inputs))
-        init_alpha0[init_residual_index, :] = 1.
+        # make sure each branch input receives from different residual stream on init
+
+        stream_branches = num_residual_streams * num_branch_inputs
+        layer_index = default(layer_index, randrange(stream_branches))
+        layer_offset = layer_index % stream_branches * num_branch_inputs
+
+        stream_seq = torch.arange(num_residual_streams)
+        branch_input_seq = torch.arange(num_branch_inputs)
+
+        init_alpha0 = rearrange(stream_seq, 's -> s 1') + rearrange(branch_input_seq, 'bi -> 1 bi') + layer_offset
+        init_alpha0 = ((init_alpha0 % num_residual_streams) == 0).float()
 
         self.static_alpha = nn.Parameter(torch.cat([init_alpha0, torch.eye(num_residual_streams)], dim = 1))
 
@@ -76,7 +83,10 @@ class HyperConnections(Module):
         self.channel_first = channel_first
 
     @classmethod
-    def get_expand_reduce_stream_functions(cls, num_streams):
+    def get_expand_reduce_stream_functions(cls, num_streams, disable = False):
+        if disable:
+            return (identity, identity)
+
         expand_fn = partial(repeat, pattern = 'b ... -> (b s) ...', s = num_streams)
         reduce_fn = partial(reduce, pattern = '(b s) ... -> b ...', reduction = 'sum', s = num_streams)
 
@@ -88,7 +98,7 @@ class HyperConnections(Module):
         hyper_conn_klass = cls if not disable else Residual
 
         init_hyper_conn_fn = partial(hyper_conn_klass, num_streams)
-        expand_reduce_fns = cls.get_expand_reduce_stream_functions(num_streams) if not disable else (identity, identity)
+        expand_reduce_fns = cls.get_expand_reduce_stream_functions(num_streams, disable = disable)
 
         return (init_hyper_conn_fn, *expand_reduce_fns)
 

{hyper_connections-0.0.14 → hyper_connections-0.0.16}/pyproject.toml

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