hyper-connections 0.4.0__tar.gz → 0.4.2__tar.gz

{hyper_connections-0.4.0 → hyper_connections-0.4.2}/PKG-INFO

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

{hyper_connections-0.4.0 → hyper_connections-0.4.2}/hyper_connections/manifold_constrained_hyper_connections.py

@@ -213,8 +213,10 @@ class ManifoldConstrainedHyperConnections(Module):
         num_fracs = 1,                      # https://arxiv.org/abs/2503.14125
         sinkhorn_iters = 20,
         log_domain_sinkhorn = False,
+        residual_mix_constraint_fn: Callable | None = None,
         forward_method_names: tuple[str, ...] = (),
-        num_dynamic_alpha_proposals = 1
+        num_dynamic_alpha_proposals = 1,
+
     ):
         """
         Appendix J, Algorithm2 in - https://arxiv.org/abs/2409.19606
@@ -287,10 +289,13 @@ class ManifoldConstrainedHyperConnections(Module):
 
             self.h_post_scale = nn.Parameter(torch.ones(()) * 1e-2)
 
-        # sinkhorn related
+        # Hres constraint related
+        # by default is sinkhorn
 
-        self.sinkhorn_iters = sinkhorn_iters
-        self.log_domain_sinkhorn = log_domain_sinkhorn
+        self.residual_mix_constraint_fn = default(
+            residual_mix_constraint_fn,
+            partial(sinkhorn_knopps if not log_domain_sinkhorn else log_domain_sinkhorn_knopps, iters = sinkhorn_iters)
+        )
 
         # dropouts
 
@@ -373,9 +378,7 @@ class ManifoldConstrainedHyperConnections(Module):
 
         alpha_pre = alpha_pre.sigmoid()
 
-        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_residual = self.residual_mix_constraint_fn(alpha_residual)
 
         alpha = cat((alpha_pre, alpha_residual), dim = -1)
 

{hyper_connections-0.4.0 → hyper_connections-0.4.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hyper-connections"
-version = "0.4.0"
+version = "0.4.2"
 description = "Hyper-Connections"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }