hyper-connections 0.3.12__py3-none-any.whl → 0.3.15__py3-none-any.whl

hyper_connections/manifold_constrained_hyper_connections.py

@@ -59,6 +59,16 @@ def sinkhorn_knopps(log_alpha, iters = 20):
 
     return alpha.to(dtype)
 
+def log_domain_sinkhorn_knopps(log_alpha, iters = 20):
+    dtype = log_alpha.dtype
+    log_alpha = log_alpha.float()
+
+    for _ in range(iters):
+        log_alpha = log_alpha - log_alpha.logsumexp(dim = -2, keepdim = True)
+        log_alpha = log_alpha - log_alpha.logsumexp(dim = -1, keepdim = True)
+
+    return log_alpha.exp()
+
 # main functions
 
 def get_expand_reduce_stream_functions(
@@ -201,6 +211,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 +283,7 @@ class ManifoldConstrainedHyperConnections(Module):
         # sinkhorn related
 
         self.sinkhorn_iters = sinkhorn_iters
+        self.log_domain_sinkhorn = log_domain_sinkhorn
 
         # dropouts
 
@@ -354,7 +366,9 @@ class ManifoldConstrainedHyperConnections(Module):
 
         alpha_pre = alpha_pre.sigmoid()
 
-        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 = sinkhorn_fn(alpha_residual, self.sinkhorn_iters)
 
         alpha = cat((alpha_pre, alpha_residual), dim = -1)
 

{hyper_connections-0.3.12.dist-info → hyper_connections-0.3.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyper-connections
-Version: 0.3.12
+Version: 0.3.15
 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.12.dist-info → hyper_connections-0.3.15.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=JF_ncPloyNLr303xZMlCiCorCE1Qt_soo1oeOyGYIdc,17168
+hyper_connections/manifold_constrained_hyper_connections.py,sha256=iS0I6Ha5iAkWdrH-dEYpyLkHRRF8yarFqHJwlTV2qLI,17689
 hyper_connections/residuals.py,sha256=JVSFJj_H7xQ3_Fd-pZH5Hdv9SveAQu29jQNvMyom5ek,921
 hyper_connections/vit.py,sha256=fTC8hAYkD4qm-KURAj8SJ66C6ZWtsBdHf_kS-4rJZGQ,5049
-hyper_connections-0.3.12.dist-info/METADATA,sha256=OyTViiZQQ1AmohzVqRhf71HOFPjLEVMGPVY5kY6wHsg,6705
-hyper_connections-0.3.12.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-hyper_connections-0.3.12.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
-hyper_connections-0.3.12.dist-info/RECORD,,
+hyper_connections-0.3.15.dist-info/METADATA,sha256=PF7nUEiHAtWHgnRQfni_eNwywX0z7_5vKyj_bNm9oi0,6705
+hyper_connections-0.3.15.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+hyper_connections-0.3.15.dist-info/licenses/LICENSE,sha256=E7RGS7kpJIStk5za_-4DVhWEAamf65EU0CNML25mq4c,1066
+hyper_connections-0.3.15.dist-info/RECORD,,