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

hyper_connections/__init__.py

@@ -1,3 +1,5 @@
 from hyper_connections.hyper_connections import (
-    HyperConnections
+    HyperConnections,
+    Residual,
+    StreamEmbed
 )

hyper_connections/hyper_connections.py

@@ -18,7 +18,46 @@ def exists(v):
 def default(v, d):
     return v if exists(v) else d
 
-# main class
+def identity(t):
+    return t
+
+# main classes
+
+# residual base class
+
+class Residual(Module):
+    def __init__(
+        self,
+        *args,
+        branch = None,
+        **kwargs
+    ):
+        super().__init__()
+        self.branch = branch
+
+    def width_connection(self, residuals):
+        return residuals, residuals, dict()
+
+    def depth_connection(self, branch_output, residuals):
+        return branch_output + residuals
+
+    def forward(self, residuals, *branch_args, **branch_kwargs):
+
+        branch_input, residuals, residual_kwargs = self.width_connection(residuals)
+
+        def add_residual_fn(branch_out):
+            return self.depth_connection(branch_out, residuals, **residual_kwargs)
+
+        if not exists(self.branch):
+            return branch_input, add_residual_fn
+
+        branch_output = self.branch(branch_input, *branch_args, **branch_kwargs)
+
+        (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 connection residual streams
 
@@ -70,10 +109,12 @@ class HyperConnections(Module):
         return expand_fn, reduce_fn
 
     @classmethod
-    def get_init_and_expand_reduce_stream_functions(cls, num_streams):
+    def get_init_and_expand_reduce_stream_functions(cls, num_streams, disable = False):
+
+        hyper_conn_klass = cls if not disable else Residual
 
-        init_hyper_conn_fn = partial(cls, num_streams)
-        expand_reduce_fns = cls.get_expand_reduce_stream_functions(num_streams)
+        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)
 
         return (init_hyper_conn_fn, *expand_reduce_fns)
 
@@ -106,9 +147,9 @@ class HyperConnections(Module):
         if self.channel_first:
             branch_input = rearrange(branch_input, 'b ... d -> b d ...')
 
-        return branch_input, residuals, beta
+        return branch_input, residuals, dict(beta = beta)
 
-    def depth_connection(self, branch_output, residuals, beta):
+    def depth_connection(self, branch_output, residuals, *, beta):
         # 'depth' connection
 
         if self.channel_first:
@@ -124,10 +165,10 @@ class HyperConnections(Module):
 
     def forward(self, residuals, *branch_args, **branch_kwargs):
 
-        branch_input, residuals, beta = self.width_connection(residuals)
+        branch_input, residuals, residual_kwargs = self.width_connection(residuals)
 
         def add_residual_fn(branch_out):
-            return self.depth_connection(branch_out, residuals, beta)
+            return self.depth_connection(branch_out, residuals, **residual_kwargs)
 
         if not exists(self.branch):
             return branch_input, add_residual_fn
@@ -139,3 +180,34 @@ class HyperConnections(Module):
         branch_output = add_residual_fn(branch_output)
 
         return tree_unflatten((branch_output, *rest), tree_spec)
+
+# stream embed
+
+class StreamEmbed(Module):
+    def __init__(
+        self,
+        num_streams,
+        dim,
+        channel_first = False
+    ):
+        super().__init__()
+        self.channel_first = channel_first
+        self.num_streams = num_streams
+
+        self.stream_embed = nn.Parameter(torch.zeros(num_streams, dim))
+
+    def forward(self, residuals):
+
+        if self.channel_first:
+            residuals = rearrange(residuals, '(b s) d ... -> b ... s d', s = self.num_streams)
+        else:
+            residuals = rearrange(residuals, '(b s) ... d -> b ... s d', s = self.num_streams)
+
+        residuals = residuals + self.stream_embed
+
+        if self.channel_first:
+            residuals = rearrange(residuals, 'b ... s d -> (b s) d ...', s = self.num_streams)
+        else:
+            residuals = rearrange(residuals, 'b ... s d -> (b s) ... d', s = self.num_streams)
+
+        return residuals

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.0.5
+Version: 0.0.7
 Summary: Hyper-Connections
 Project-URL: Homepage, https://pypi.org/project/hyper-connections/
 Project-URL: Repository, https://github.com/lucidrains/hyper-connections
@@ -135,6 +135,12 @@ residual = add_residual(branch_output)
 residual = reduce_stream(residual)
 ```
 
+To compare hyper connections to plain residual without changing the code, just pass `disable = True` when fetching the functions
+
+```python
+HyperConnections.get_init_and_expand_reduce_stream_functions(4, disable = True)
+```
+
 ## Citation
 
 ```bibtex

hyper_connections-0.0.7.dist-info/RECORD

@@ -0,0 +1,6 @@
+hyper_connections/__init__.py,sha256=mAy66IuHqXM4XOyOZGt5mo2B2hfHdUk8jW31YnWNQTg,104
+hyper_connections/hyper_connections.py,sha256=RBm0qEhQwCSlvtqNXo_YIkRAkMynCNBrN7xXt4rsRBc,6756
+hyper_connections-0.0.7.dist-info/METADATA,sha256=PTozroByBHtvwj8fFJhflo0H0GdwRUM__8aNP2LzuPY,4978
+hyper_connections-0.0.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+hyper_connections-0.0.7.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
+hyper_connections-0.0.7.dist-info/RECORD,,

hyper_connections-0.0.5.dist-info/RECORD

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