Make project_weights public

Author: PierreQuinton

src/torchjd/aggregation/_dual_cone_utils.py

@@ -6,7 +6,7 @@
 from torch import Tensor
 
 
-def _project_weights(U: Tensor, G: Tensor, solver: Literal["quadprog"]) -> Tensor:
+def project_weights(U: Tensor, G: Tensor, solver: Literal["quadprog"]) -> Tensor:
     """
     Computes the tensor of weights corresponding to the projection of the vectors in `U` onto the
     rows of a matrix whose Gramian is provided.

src/torchjd/aggregation/dualproj.py

@@ -2,7 +2,7 @@
 
 from torch import Tensor
 
-from ._dual_cone_utils import _project_weights
+from ._dual_cone_utils import project_weights
 from ._gramian_utils import _compute_regularized_normalized_gramian
 from ._pref_vector_utils import _pref_vector_to_str_suffix, _pref_vector_to_weighting
 from .bases import _WeightedAggregator, _Weighting
@@ -101,5 +101,5 @@ def __init__(
     def forward(self, matrix: Tensor) -> Tensor:
         u = self.weighting(matrix)
         G = _compute_regularized_normalized_gramian(matrix, self.norm_eps, self.reg_eps)
-        w = _project_weights(u, G, self.solver)
+        w = project_weights(u, G, self.solver)
         return w

src/torchjd/aggregation/upgrad.py

@@ -3,7 +3,7 @@
 import torch
 from torch import Tensor
 
-from ._dual_cone_utils import _project_weights
+from ._dual_cone_utils import project_weights
 from ._gramian_utils import _compute_regularized_normalized_gramian
 from ._pref_vector_utils import _pref_vector_to_str_suffix, _pref_vector_to_weighting
 from .bases import _WeightedAggregator, _Weighting
@@ -97,5 +97,5 @@ def __init__(
     def forward(self, matrix: Tensor) -> Tensor:
         U = torch.diag(self.weighting(matrix))
         G = _compute_regularized_normalized_gramian(matrix, self.norm_eps, self.reg_eps)
-        W = _project_weights(U, G, self.solver)
+        W = project_weights(U, G, self.solver)
         return torch.sum(W, dim=0)

tests/unit/aggregation/test_dual_cone_utils.py

@@ -3,7 +3,7 @@
 from pytest import mark, raises
 from torch.testing import assert_close
 
-from torchjd.aggregation._dual_cone_utils import _project_weight_vector, _project_weights
+from torchjd.aggregation._dual_cone_utils import _project_weight_vector, project_weights
 
 
 @mark.parametrize("shape", [(5, 7), (9, 37), (2, 14), (32, 114), (50, 100)])
@@ -33,7 +33,7 @@ def test_solution_weights(shape: tuple[int, int]):
     G = J @ J.T
     u = torch.rand(shape[0])
 
-    w = _project_weights(u, G, "quadprog")
+    w = project_weights(u, G, "quadprog")
     dual_gap = w - u
 
     # Dual feasibility
@@ -64,8 +64,8 @@ def test_tensorization_shape(shape: tuple[int, ...]):
 
     G = matrix @ matrix.T
 
-    W_tensor = _project_weights(U_tensor, G, "quadprog")
-    W_matrix = _project_weights(U_matrix, G, "quadprog")
+    W_tensor = project_weights(U_tensor, G, "quadprog")
+    W_matrix = project_weights(U_matrix, G, "quadprog")
 
     assert_close(W_matrix.reshape(shape), W_tensor)