hyper-connections 0.0.6__tar.gz → 0.0.7__tar.gz

{hyper_connections-0.0.6 → hyper_connections-0.0.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.0.6
+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
@@ -138,7 +138,7 @@ 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, disabled = True)
+HyperConnections.get_init_and_expand_reduce_stream_functions(4, disable = True)
 ```
 
 ## Citation

{hyper_connections-0.0.6 → hyper_connections-0.0.7}/README.md

@@ -95,7 +95,7 @@ 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, disabled = True)
+HyperConnections.get_init_and_expand_reduce_stream_functions(4, disable = True)
 ```
 
 ## Citation

{hyper_connections-0.0.6 → hyper_connections-0.0.7}/hyper_connections/__init__.py

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

{hyper_connections-0.0.6 → hyper_connections-0.0.7}/hyper_connections/hyper_connections.py

@@ -180,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.6 → hyper_connections-0.0.7}/pyproject.toml

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