hyper-connections 0.4.7__py3-none-any.whl → 0.4.8__py3-none-any.whl

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,6 +383,14 @@ class ManifoldConstrainedHyperConnections(Module):
 
         residuals = self.split_fracs(residuals)
 
+        # maybe mix streams
+
+        if exists(self.maybe_mix_streams):
+
+            residuals, inverse_pack_lead_dims = pack_with_inverse(residuals, '* c h w')
+            residuals = self.maybe_mix_streams(residuals)
+            residuals = inverse_pack_lead_dims(residuals)
+
         # norm
 
         normed = self.norm(residuals)

{hyper_connections-0.4.7.dist-info → hyper_connections-0.4.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.4.7
+Version: 0.4.8
 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.dist-info → hyper_connections-0.4.8.dist-info}/RECORD

@@ -3,12 +3,12 @@ hyper_connections/hyper_connections.py,sha256=2F-104cGE82KCK0KeC07NSOJNPT-0PCtvX
 hyper_connections/hyper_connections_channel_first.py,sha256=5vAen4WXxNI9K07ndLBQJwdJv-OjoXznta5EIQTaQNA,6512
 hyper_connections/hyper_connections_with_multi_branch_inputs.py,sha256=yn2AlFB6qCYQeRhhhaMlCM3mxxLEdWCYwU2p9TsMwWI,7835
 hyper_connections/hyper_connections_with_multi_input_streams.py,sha256=iFPw5pgCRHTo16nBJ2PExKSrvTyCh7ba7Py14P1oSPE,11311
-hyper_connections/mHCv2.py,sha256=XB2HwxTo7daZvy9fzF8UjHI12ephwgE91h9AH2Ou4WI,17452
+hyper_connections/mHCv2.py,sha256=LpMtlrb7Vfi2qq_cqPl9fajA5SxkMTl5QGpmvBJyD1M,18360
 hyper_connections/manifold_constrained_hyper_connections.py,sha256=E4os-6q_SMjJO1JD0EG8rFTCXA7MQoy-aqUlM7KVS5Q,18269
 hyper_connections/residuals.py,sha256=JVSFJj_H7xQ3_Fd-pZH5Hdv9SveAQu29jQNvMyom5ek,921
 hyper_connections/triton_sinkhorn.py,sha256=n2WyQcUemtv5T5Sk2nljnSpV2hEED4I3HaPsIUy4638,5905
 hyper_connections/vit.py,sha256=dh8AVMUPaUHuWxXJEHoMW_G5nj-EQQjDmgbPwwhiq5g,5215
-hyper_connections-0.4.7.dist-info/METADATA,sha256=2ajn-IuCxuUjgnOw5dEBxXKqLJbyQohHSGgKJ2dZFoA,6704
-hyper_connections-0.4.7.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-hyper_connections-0.4.7.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
-hyper_connections-0.4.7.dist-info/RECORD,,
+hyper_connections-0.4.8.dist-info/METADATA,sha256=vevhBHad-7ffu1KBFcazUqU5C2XVRy1LlZkIxJUNDIs,6746
+hyper_connections-0.4.8.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+hyper_connections-0.4.8.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
+hyper_connections-0.4.8.dist-info/RECORD,,