hyper-connections 0.0.17__tar.gz → 0.0.19__tar.gz

{hyper_connections-0.0.17 → hyper_connections-0.0.19}/PKG-INFO

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

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

{hyper_connections-0.0.17 → hyper_connections-0.0.19}/hyper_connections/hyper_connections.py

@@ -86,7 +86,8 @@ class HyperConnections(Module):
         branch: Module | None = None,
         layer_index = None,
         tanh = True,
-        channel_first = False
+        channel_first = False,
+        dropout = 0.
     ):
         """
         Appendix J, Algorithm2 in - https://arxiv.org/abs/2409.19606
@@ -116,6 +117,10 @@ class HyperConnections(Module):
         self.dynamic_beta_fn = nn.Parameter(torch.zeros(dim))
         self.dynamic_beta_scale = nn.Parameter(torch.ones(()) * 1e-2)
 
+        # dropouts
+
+        self.dropout = nn.Dropout(dropout)
+
         # channel first option
 
         self.channel_first = channel_first
@@ -184,7 +189,7 @@ class HyperConnections(Module):
         if self.channel_first:
             output = rearrange(output, 'b ... d -> b d ...')
 
-        return output
+        return self.dropout(output)
 
     def decorate_branch(self, branch: Callable):
         assert not exists(self.branch), 'branch was already wrapped on init'
@@ -253,3 +258,36 @@ class StreamEmbed(Module):
             residuals = rearrange(residuals, 'b ... s d -> (b s) ... d', s = self.num_streams)
 
         return residuals
+
+# attention pool - taken from Enformer https://www.nature.com/articles/s41592-021-01252-x , in turn taken from somewhere else
+
+class AttentionPoolReduceStream(Module):
+    def __init__(
+        self,
+        num_streams,
+        dim,
+        channel_first = False
+    ):
+        super().__init__()
+        self.num_streams = num_streams
+        self.channel_first = channel_first
+
+        self.to_attn_logits = nn.Linear(dim, dim, bias = False)
+        self.to_attn_logits.weight.data.copy_(torch.eye(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)
+
+        attn_logits = self.to_attn_logits(residuals)
+        attn = attn_logits.softmax(dim = -2)
+
+        residuals = reduce(residuals * attn, 'b ... s d -> b ... d', 'sum')
+
+        if self.channel_first:
+            residuals = rearrange(residuals, 'b ... d -> b d ...')
+
+        return residuals

{hyper_connections-0.0.17 → hyper_connections-0.0.19}/hyper_connections/hyper_connections_with_multi_branch_inputs.py

@@ -42,6 +42,7 @@ class HyperConnections(Module):
         layer_index = None,
         tanh = True,
         channel_first = False,
+        dropout = 0.,
         num_branch_inputs = 1  # residuals will be linearly combined to multiple inputs, fed through the branch, then linearly combined back out to residuals
     ):
         """
@@ -89,6 +90,10 @@ class HyperConnections(Module):
         self.dynamic_beta_fn = nn.Parameter(torch.zeros(dim, num_branch_inputs))
         self.dynamic_beta_scale = nn.Parameter(torch.ones(()) * 1e-2)
 
+        # dropout
+
+        self.dropout = nn.Dropout(dropout)
+
         # channel first option
 
         self.channel_first = channel_first
@@ -164,7 +169,7 @@ class HyperConnections(Module):
         if self.channel_first:
             output = rearrange(output, 'b ... d -> b d ...')
 
-        return output
+        return self.dropout(output)
 
     def decorate_branch(self, branch: Callable | tuple[Callable, ...] | list[Callable]):
         assert not exists(self.branches), 'branch was already wrapped on init'

{hyper_connections-0.0.17 → hyper_connections-0.0.19}/pyproject.toml

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