hyper-connections 0.3.11__py3-none-any.whl → 0.3.14__py3-none-any.whl

hyper_connections/manifold_constrained_hyper_connections.py

@@ -59,6 +59,14 @@ def sinkhorn_knopps(log_alpha, iters = 20):
 
     return alpha.to(dtype)
 
+def log_domain_sinkhorn_knopps(alpha, iters = 20):
+
+    for _ in range(iters):
+        alpha = alpha - alpha.logsumexp(dim = -2, keepdim = True)
+        alpha = alpha - alpha.logsumexp(dim = -1, keepdim = True)
+
+    return alpha.exp()
+
 # main functions
 
 def get_expand_reduce_stream_functions(
@@ -201,6 +209,7 @@ class ManifoldConstrainedHyperConnections(Module):
         depth_residual_fn = add,
         num_fracs = 1,                      # https://arxiv.org/abs/2503.14125
         sinkhorn_iters = 20,
+        log_domain_sinkhorn = False,
         forward_method_names: tuple[str, ...] = (),
     ):
         """
@@ -272,6 +281,7 @@ class ManifoldConstrainedHyperConnections(Module):
         # sinkhorn related
 
         self.sinkhorn_iters = sinkhorn_iters
+        self.log_domain_sinkhorn = log_domain_sinkhorn
 
         # dropouts
 
@@ -348,22 +358,20 @@ class ManifoldConstrainedHyperConnections(Module):
 
         alpha = dynamic_alpha + static_alpha
 
-        alpha = self.split_fracs(alpha) # (batch, seq, fracs1, streams, fracs2, input + residual streams)
-
         # the alpha is now split and "manifold constrained" with sinkhorn and sigmoid
 
         alpha_pre, alpha_residual = alpha[..., :self.num_input_views], alpha[..., self.num_input_views:]
 
         alpha_pre = alpha_pre.sigmoid()
 
-        alpha_residual = rearrange(alpha_residual, '... (v s1 s2) -> ... v s1 s2', v = self.num_input_views, s1 = streams)
-
-        alpha_residual = sinkhorn_knopps(alpha_residual, self.sinkhorn_iters)
+        sinkhorn_fn = sinkhorn_knopps if not self.log_domain_sinkhorn else log_domain_sinkhorn_knopps
 
-        alpha_residual = rearrange(alpha_residual, '... v s1 s2 -> ... (v s1 s2)')
+        alpha_residual = sinkhorn_fn(alpha_residual, self.sinkhorn_iters)
 
         alpha = cat((alpha_pre, alpha_residual), dim = -1)
 
+        alpha = self.split_fracs(alpha) # (batch, seq, fracs1, streams, fracs2, input + residual streams)
+
         # beta for weights from branch output back to residual streams
 
         beta = None

{hyper_connections-0.3.11.dist-info → hyper_connections-0.3.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.3.11
+Version: 0.3.14
 Summary: Hyper-Connections
 Project-URL: Homepage, https://pypi.org/project/hyper-connections/
 Project-URL: Repository, https://github.com/lucidrains/hyper-connections

{hyper_connections-0.3.11.dist-info → hyper_connections-0.3.14.dist-info}/RECORD

@@ -3,10 +3,10 @@ hyper_connections/hyper_connections.py,sha256=rqFJj3U0LF3uDKNKNPBpRrmf0oa2BGWVbD
 hyper_connections/hyper_connections_channel_first.py,sha256=Mh_hzhTi96ZoOPmhSKwUaF4TbHpNqhs83wNe5hNuL7o,6532
 hyper_connections/hyper_connections_with_multi_branch_inputs.py,sha256=6BXKdSwyx6wdQVseebKG2EQkhVaVLrrepOlL8lLnex4,7855
 hyper_connections/hyper_connections_with_multi_input_streams.py,sha256=ueT3CJPHrt5hRU7q1bFF0rANWJh_pXqclt6HiUu1gBY,11331
-hyper_connections/manifold_constrained_hyper_connections.py,sha256=u_yeGVdPCMvZnPzj3cM_Kfo0XoZRChU4z-qL8oIKnRQ,17376
+hyper_connections/manifold_constrained_hyper_connections.py,sha256=BekY37Gt5us7zZjxP6laDHCbygMww8kS5GatnveNrYw,17595
 hyper_connections/residuals.py,sha256=JVSFJj_H7xQ3_Fd-pZH5Hdv9SveAQu29jQNvMyom5ek,921
 hyper_connections/vit.py,sha256=fTC8hAYkD4qm-KURAj8SJ66C6ZWtsBdHf_kS-4rJZGQ,5049
-hyper_connections-0.3.11.dist-info/METADATA,sha256=HPZFYGNNJENlhuSEdKAdUUPEBaQ88ikkwViItlZtBHs,6705
-hyper_connections-0.3.11.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-hyper_connections-0.3.11.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
-hyper_connections-0.3.11.dist-info/RECORD,,
+hyper_connections-0.3.14.dist-info/METADATA,sha256=xncYFa2ttI2WOVLTy5-jk69dZ29aQa-QfFAgI4qJvHc,6705
+hyper_connections-0.3.14.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+hyper_connections-0.3.14.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
+hyper_connections-0.3.14.dist-info/RECORD,,