hyper-connections 0.4.7__tar.gz → 0.4.9__tar.gz

{hyper_connections-0.4.7 → hyper_connections-0.4.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.4.7
+Version: 0.4.9
 Summary: Hyper-Connections
 Project-URL: Homepage, https://pypi.org/project/hyper-connections/
 Project-URL: Repository, https://github.com/lucidrains/hyper-connections
@@ -35,6 +35,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Requires-Dist: einops>=0.8.1
+Requires-Dist: torch-einops-utils>=0.0.20
 Requires-Dist: torch>=2.5
 Provides-Extra: examples
 Description-Content-Type: text/markdown

{hyper_connections-0.4.7 → hyper_connections-0.4.9}/hyper_connections/mHCv2.py

@@ -13,6 +13,8 @@ from torch.utils._pytree import tree_flatten, tree_unflatten
 from einops import rearrange, repeat, reduce, einsum
 from einops.layers.torch import Rearrange, Reduce
 
+from torch_einops_utils import pack_with_inverse
+
 """
 ein notation:
 b - batch
@@ -241,6 +243,7 @@ class ManifoldConstrainedHyperConnections(Module):
         forward_method_names: tuple[str, ...] = (),
         num_dynamic_alpha_proposals = 1,
         use_triton_sinkhorn = False,
+        mix_streams_before_norm = False, # whether to mix the residual streams before the norm (that then projects to Hpre, Hpost, Hresidual)
     ):
         """
         Appendix J, Algorithm2 in - https://arxiv.org/abs/2409.19606
@@ -263,6 +266,16 @@ class ManifoldConstrainedHyperConnections(Module):
 
         dim //= num_fracs # effective dim handled in dimension is feature dimension divided by num fractions
 
+        # whether to mix the streams before the norm below
+        # this would be equivalent to separable depthwise convs from yesteryears (with a norm in between) - parameter efficient improv
+
+        self.maybe_mix_streams = None
+
+        if mix_streams_before_norm:
+            self.maybe_mix_streams = nn.Conv2d(num_residual_streams, num_residual_streams, 1, bias = False)
+
+            nn.init.dirac_(self.maybe_mix_streams.weight)
+
         # they used layernorm in paper, but rmsnorm is fine given what we know now
 
         self.norm = RMSNorm(dim)
@@ -370,9 +383,19 @@ class ManifoldConstrainedHyperConnections(Module):
 
         residuals = self.split_fracs(residuals)
 
+        # maybe mix streams
+
+        norm_input = residuals
+
+        if exists(self.maybe_mix_streams):
+
+            norm_input, inverse_pack_lead_dims = pack_with_inverse(norm_input, '* c h w')
+            norm_input = self.maybe_mix_streams(norm_input)
+            norm_input = inverse_pack_lead_dims(norm_input)
+
         # norm
 
-        normed = self.norm(residuals)
+        normed = self.norm(norm_input)
 
         # alpha for weighted sum of residuals going into branch
 

{hyper_connections-0.4.7 → hyper_connections-0.4.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hyper-connections"
-version = "0.4.7"
+version = "0.4.9"
 description = "Hyper-Connections"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -25,6 +25,7 @@ classifiers=[
 dependencies = [
     "einops>=0.8.1",
     "torch>=2.5",
+    "torch-einops-utils>=0.0.20"
 ]
 
 [project.urls]

{hyper_connections-0.4.7 → hyper_connections-0.4.9}/tests/test_hyper_connections.py

@@ -227,7 +227,9 @@ def test_mhc_vit(
         dropout = 0.1,
         emb_dropout = 0.1,
         num_residual_streams = 4,
-        num_dynamic_alpha_proposals = num_dynamic_alpha_proposals
+        mhc_kwargs = dict(
+            num_dynamic_alpha_proposals = num_dynamic_alpha_proposals
+        )
     )
 
     img = torch.randn(1, 3, 256, 256)
@@ -238,11 +240,13 @@ def test_mhc_vit(
 @param('num_fracs', (1, 2))
 @param('num_streams', (1, 3))
 @param('disable', (False, True))
+@param('mix_streams_before_norm', (False, True))
 @param('add_attn_pool_reduce_stream', (False, True))
 def test_mhcv2(
     num_fracs,
     num_streams,
     disable,
+    mix_streams_before_norm,
     add_attn_pool_reduce_stream
 ):
     import torch
@@ -261,7 +265,7 @@ def test_mhcv2(
 
     from hyper_connections.mHCv2 import get_init_and_expand_reduce_stream_functions
 
-    init_hyper_conn, expand_stream, reduce_stream = get_init_and_expand_reduce_stream_functions(num_streams, dim = 512, num_fracs = num_fracs, disable = disable, add_attn_pool_reduce_stream = add_attn_pool_reduce_stream)
+    init_hyper_conn, expand_stream, reduce_stream = get_init_and_expand_reduce_stream_functions(num_streams, dim = 512, num_fracs = num_fracs, mix_streams_before_norm = mix_streams_before_norm, disable = disable, add_attn_pool_reduce_stream = add_attn_pool_reduce_stream)
 
     # 1. wrap your branch function