hyper-connections 0.0.15__py3-none-any.whl → 0.0.17__py3-none-any.whl

hyper_connections/hyper_connections_with_multi_branch_inputs.py

@@ -6,8 +6,8 @@ from random import randrange
 
 import torch
 from torch import nn
-from torch.nn import Module
 import torch.nn.functional as F
+from torch.nn import Module, ModuleList
 from torch.utils._pytree import tree_flatten, tree_unflatten
 
 from einops import rearrange, repeat, reduce, einsum
@@ -22,6 +22,9 @@ def exists(v):
 def default(v, d):
     return v if exists(v) else d
 
+def divisible_by(num, den):
+    return (num % den) == 0
+
 def identity(t):
     return t
 
@@ -35,7 +38,7 @@ class HyperConnections(Module):
         num_residual_streams,
         *,
         dim,
-        branch: Module | None = None,
+        branch: Module | tuple[Module, ...] | list[Module] | None = None,
         layer_index = None,
         tanh = True,
         channel_first = False,
@@ -46,7 +49,15 @@ class HyperConnections(Module):
         """
         super().__init__()
 
-        self.branch = branch
+        self.branches = None
+
+        if isinstance(branch, Module):
+            branch = [branch]
+
+        if exists(branch):
+            assert divisible_by(num_branch_inputs, len(branch))
+
+            self.branches = ModuleList(branch)
 
         # activation, seemingly results were wishy washy depending on using tanh or not
 
@@ -57,12 +68,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 +94,7 @@ 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)
 
@@ -148,13 +166,22 @@ class HyperConnections(Module):
 
         return output
 
-    def decorate_branch(self, branch: Callable):
-        assert not exists(self.branch), 'branch was already wrapped on init'
+    def decorate_branch(self, branch: Callable | tuple[Callable, ...] | list[Callable]):
+        assert not exists(self.branches), 'branch was already wrapped on init'
 
         def forward_and_add_residual(residual, *args, **kwargs):
             branch_input, add_residual = self.forward(residual)
 
-            branch_output = branch(branch_input, *args, **kwargs)
+            if callable(branch):
+                branches = [branch]
+            else:
+                branches = branch
+
+            branch_inputs = rearrange(branch_input, '(br b) ... -> br b ...', br = len(branches))
+
+            branch_outputs = [fn(x, *args, **kwargs) for fn, x in zip(branches, branch_inputs)]
+
+            branch_output = torch.cat(branch_outputs)
 
             residual = add_residual(branch_output)
 
@@ -173,9 +200,13 @@ class HyperConnections(Module):
 
             return tree_unflatten((branch_out, *rest), tree_spec)
 
-        if not exists(self.branch):
+        if not exists(self.branches):
             return branch_input, add_residual_fn
 
-        branch_output = self.branch(branch_input, *branch_args, **branch_kwargs)
+        branch_inputs = rearrange(branch_input, '(br b) ... -> br b ...', br = len(self.branches))
+
+        branch_outputs = [fn(x, *branch_args, **branch_kwargs) for fn, x in zip(self.branches, branch_inputs)]
+
+        branch_output = torch.cat(branch_outputs)
 
         return add_residual_fn(branch_output)

{hyper_connections-0.0.15.dist-info → hyper_connections-0.0.17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.0.15
+Version: 0.0.17
 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.17.dist-info/RECORD

@@ -0,0 +1,7 @@
+hyper_connections/__init__.py,sha256=mAy66IuHqXM4XOyOZGt5mo2B2hfHdUk8jW31YnWNQTg,104
+hyper_connections/hyper_connections.py,sha256=LoyiGGDH81MLUcF9WKT39lA9Y1CXUAd-SYM9tJGQ61A,8099
+hyper_connections/hyper_connections_with_multi_branch_inputs.py,sha256=l8oKkeVouFWQZY_0IVNTtgQPSQTA9qZdgLrfotRS_5w,7297
+hyper_connections-0.0.17.dist-info/METADATA,sha256=dB2ZvPv63hTNJB6FEbyLlqxer4L7epwFMOvEVNjNOGk,5076
+hyper_connections-0.0.17.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+hyper_connections-0.0.17.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
+hyper_connections-0.0.17.dist-info/RECORD,,

hyper_connections-0.0.15.dist-info/RECORD

@@ -1,7 +0,0 @@
-hyper_connections/__init__.py,sha256=mAy66IuHqXM4XOyOZGt5mo2B2hfHdUk8jW31YnWNQTg,104
-hyper_connections/hyper_connections.py,sha256=LoyiGGDH81MLUcF9WKT39lA9Y1CXUAd-SYM9tJGQ61A,8099
-hyper_connections/hyper_connections_with_multi_branch_inputs.py,sha256=6KHdqvjUZCyHwf7LVByMMuoQtnEHBxRxdt-4Bw9HPcA,6130
-hyper_connections-0.0.15.dist-info/METADATA,sha256=DnZongXit3VTaT7DxmpK0RRfUhXqfUTjM32a-ofbrIw,5076
-hyper_connections-0.0.15.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-hyper_connections-0.0.15.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
-hyper_connections-0.0.15.dist-info/RECORD,,