hyper-connections 0.1.9__py3-none-any.whl → 0.1.11__py3-none-any.whl

hyper_connections/hyper_connections.py

@@ -19,6 +19,7 @@ b - batch
 d - feature dimension
 s - residual streams
 t - residual streams + num branch inputs
+v - number of views for branch input
 """
 
 # helper functions
@@ -34,22 +35,31 @@ def identity(t):
 
 # main functions
 
-def get_expand_reduce_stream_functions(num_streams, disable = False):
+def get_expand_reduce_stream_functions(num_streams, add_stream_embed = False, dim = None, disable = False):
 
     if num_streams == 1 or disable:
         return (nn.Identity(), nn.Identity())
 
-    expand_fn = Reduce(pattern = 'b ... -> (b s) ...', reduction = 'repeat', s = num_streams)
+    if add_stream_embed:
+        assert exists(dim), '`dim` must be passed into get_init_and_expand_reduce_stream_functions for returning an expansion function with stream embeddings added'
+
+        expand_fn = StreamEmbed(num_streams, dim, expand_to_streams = True)
+    else:
+        expand_fn = Reduce(pattern = 'b ... -> (b s) ...', reduction = 'repeat', s = num_streams)
+
     reduce_fn = Reduce(pattern = '(b s) ... -> b ...', reduction = 'sum', s = num_streams)
 
     return expand_fn, reduce_fn
 
-def get_init_and_expand_reduce_stream_functions(num_streams, disable = False):
+def get_init_and_expand_reduce_stream_functions(num_streams, dim = None, add_stream_embed = False, disable = False):
 
     hyper_conn_klass = HyperConnections if not disable else Residual
 
     init_hyper_conn_fn = partial(hyper_conn_klass, num_streams)
-    expand_reduce_fns = get_expand_reduce_stream_functions(num_streams, disable = disable)
+    expand_reduce_fns = get_expand_reduce_stream_functions(num_streams, add_stream_embed = add_stream_embed, dim = dim, disable = disable)
+
+    if exists(dim):
+        init_hyper_conn_fn = partial(init_hyper_conn_fn, dim = dim)
 
     return (init_hyper_conn_fn, *expand_reduce_fns)
 
@@ -132,7 +142,8 @@ class HyperConnections(Module):
         channel_first = False,
         dropout = 0.,
         residual_transform: Module | None = None, # to support resnet blocks where dimension in not equal to dimension out - usually a residual conv
-        add_branch_out_to_residual = True # will disable depth connections (weighted residual sum with beta) if set False
+        add_branch_out_to_residual = True, # will disable depth connections (weighted residual sum with beta) if set False
+        num_input_views = 1 # allow for the branch module to receive multiple input views, dimension placed on the very left (before batch)
     ):
         """
         Appendix J, Algorithm2 in - https://arxiv.org/abs/2409.19606
@@ -152,14 +163,19 @@ class HyperConnections(Module):
         self.num_residual_streams = num_residual_streams
         init_residual_index = default(layer_index, randrange(num_residual_streams)) % num_residual_streams # just choose one random residual stream if layer index not given
 
+        # width num residual streams
+
+        assert num_input_views >= 1
+        self.num_input_views = num_input_views
+
         # width connection
 
-        init_alpha0 = torch.zeros((num_residual_streams, 1))
-        init_alpha0[init_residual_index, 0] = 1.
+        init_alpha0 = torch.zeros((num_residual_streams, num_input_views))
+        init_alpha0[init_residual_index, :] = 1.
 
         self.static_alpha = nn.Parameter(torch.cat([init_alpha0, torch.eye(num_residual_streams)], dim = 1))
 
-        self.dynamic_alpha_fn = nn.Parameter(torch.zeros(dim, num_residual_streams + 1))
+        self.dynamic_alpha_fn = nn.Parameter(torch.zeros(dim, num_residual_streams + num_input_views))
         self.dynamic_alpha_scale = nn.Parameter(torch.ones(()) * 1e-2)
 
         # depth connection related (beta)
@@ -213,7 +229,11 @@ class HyperConnections(Module):
 
         mix_h = einsum(alpha, residuals, '... s t, ... s d -> ... t d')
 
-        branch_input, residuals = mix_h[..., 0, :], mix_h[..., 1:, :]
+        if self.num_input_views == 1:
+            branch_input, residuals = mix_h[..., 0, :], mix_h[..., 1:, :]
+        else:
+            branch_input, residuals = mix_h[..., :self.num_input_views, :], mix_h[..., self.num_input_views:, :]
+            branch_input = rearrange(branch_input, 'b ... v d -> v b ... d')
 
         if self.channel_first:
             branch_input = rearrange(branch_input, 'b ... d -> b d ...')

{hyper_connections-0.1.9.dist-info → hyper_connections-0.1.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.1.9
+Version: 0.1.11
 Summary: Hyper-Connections
 Project-URL: Homepage, https://pypi.org/project/hyper-connections/
 Project-URL: Repository, https://github.com/lucidrains/hyper-connections

hyper_connections-0.1.11.dist-info/RECORD

@@ -0,0 +1,9 @@
+hyper_connections/__init__.py,sha256=d2zNTka0Gp9vINu4U-RhgTJFBhsVrs1fne_15Zl0oOs,224
+hyper_connections/hyper_connections.py,sha256=mUImPtaTE8Paygs-6vq7l_mlph1CkU__jRcE4TFim_Y,12137
+hyper_connections/hyper_connections_channel_first.py,sha256=BojfO2dcT4jX1rlcU3kr0B6B_CjrkkS2AZU4ZXeWvh8,6769
+hyper_connections/hyper_connections_with_multi_branch_inputs.py,sha256=2JABz1slrF5_XP33L1CMNTmmixfoD464FtQpUADqneU,7806
+hyper_connections/hyper_connections_with_multi_input_streams.py,sha256=UIKceEZEBLKFL5VuP5tR1KTDeZNIJEKjFuPAwXkcp0I,11282
+hyper_connections-0.1.11.dist-info/METADATA,sha256=Ck3udilJMrT1ABRkqkNhfEdkjQXAKFBtsUAAalEu3No,5231
+hyper_connections-0.1.11.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+hyper_connections-0.1.11.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
+hyper_connections-0.1.11.dist-info/RECORD,,

hyper_connections-0.1.9.dist-info/RECORD

@@ -1,9 +0,0 @@
-hyper_connections/__init__.py,sha256=d2zNTka0Gp9vINu4U-RhgTJFBhsVrs1fne_15Zl0oOs,224
-hyper_connections/hyper_connections.py,sha256=F81iJkGMpxgCZPaBTLf0c3CYIE-ROAVgZJWY3NlrsJw,11068
-hyper_connections/hyper_connections_channel_first.py,sha256=BojfO2dcT4jX1rlcU3kr0B6B_CjrkkS2AZU4ZXeWvh8,6769
-hyper_connections/hyper_connections_with_multi_branch_inputs.py,sha256=2JABz1slrF5_XP33L1CMNTmmixfoD464FtQpUADqneU,7806
-hyper_connections/hyper_connections_with_multi_input_streams.py,sha256=UIKceEZEBLKFL5VuP5tR1KTDeZNIJEKjFuPAwXkcp0I,11282
-hyper_connections-0.1.9.dist-info/METADATA,sha256=XxicphOwzNfTmBLF4Py89MhTWSpJqVk1EG-DV1gpFvo,5230
-hyper_connections-0.1.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-hyper_connections-0.1.9.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
-hyper_connections-0.1.9.dist-info/RECORD,,