hyper-connections 0.0.18__py3-none-any.whl → 0.0.20__py3-none-any.whl

hyper_connections/__init__.py

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

hyper_connections/hyper_connections.py

@@ -12,6 +12,14 @@ from torch.utils._pytree import tree_flatten, tree_unflatten
 
 from einops import rearrange, repeat, reduce, einsum
 
+"""
+ein notation:
+b - batch
+d - feature dimension
+s - residual streams
+t - residual streams + num branch inputs
+"""
+
 # helper functions
 
 def exists(v):
@@ -102,6 +110,8 @@ class HyperConnections(Module):
 
         self.norm = nn.RMSNorm(dim) # they used layernorm in paper, but rmsnorm is fine given what we know now
 
+        assert num_residual_streams > 0, '`num_residual_streams` must be greater than 0'
+
         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
 
@@ -258,3 +268,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/hyper_connections_with_multi_branch_inputs.py

@@ -12,6 +12,16 @@ from torch.utils._pytree import tree_flatten, tree_unflatten
 
 from einops import rearrange, repeat, reduce, einsum
 
+"""
+ein notation:
+b - batch
+d - feature dimension
+s - residual streams
+i - branch inputs
+br - branch functions
+t - residual streams + num branch inputs
+"""
+
 from hyper_connections.hyper_connections import Residual, StreamEmbed
 
 # helper functions

{hyper_connections-0.0.18.dist-info → hyper_connections-0.0.20.dist-info}/METADATA

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

@@ -0,0 +1,7 @@
+hyper_connections/__init__.py,sha256=wJxbrEXRGmOIjPw8fWP-cUq6CE8bvx95mIlhWifNvYc,135
+hyper_connections/hyper_connections.py,sha256=ElPtieRLvVKaVg2Attx1k6esKq1SY2X4AVZbZmsQAOM,9486
+hyper_connections/hyper_connections_with_multi_branch_inputs.py,sha256=HbLpt79xcMv_os6brMvDd90t2GOPceliE1YFusR2eJI,7553
+hyper_connections-0.0.20.dist-info/METADATA,sha256=erA-d7KNNdzPY76x8IWKd2trv2WuBO9-C2DtH-SoQ_Y,5076
+hyper_connections-0.0.20.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+hyper_connections-0.0.20.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
+hyper_connections-0.0.20.dist-info/RECORD,,

hyper_connections-0.0.18.dist-info/RECORD

@@ -1,7 +0,0 @@
-hyper_connections/__init__.py,sha256=mAy66IuHqXM4XOyOZGt5mo2B2hfHdUk8jW31YnWNQTg,104
-hyper_connections/hyper_connections.py,sha256=hZ-O79zKOAJxAzobAQVHGamQUEnZbs5tD2vqonATLUY,8199
-hyper_connections/hyper_connections_with_multi_branch_inputs.py,sha256=sTQ4sh1JLhRz06iP4PvpkLb_BMvIrcWRC-5_JxkbznQ,7396
-hyper_connections-0.0.18.dist-info/METADATA,sha256=GzMgZhbW5wLLHTsTbAU2lk1-8cGL9UNBt8iejNr08xA,5076
-hyper_connections-0.0.18.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-hyper_connections-0.0.18.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
-hyper_connections-0.0.18.dist-info/RECORD,,