Make functions in _gramian_utils public (for the same package)

Author: PierreQuinton

src/torchjd/aggregation/_gramian_utils.py

@@ -2,15 +2,15 @@
 from torch import Tensor
 
 
-def _compute_gramian(matrix: Tensor) -> Tensor:
+def compute_gramian(matrix: Tensor) -> Tensor:
     """
     Computes the `Gramian matrix <https://en.wikipedia.org/wiki/Gram_matrix>`_ of a given matrix.
     """
 
     return matrix @ matrix.T
 
 
-def _normalize(gramian: Tensor, eps: float) -> Tensor:
+def normalize(gramian: Tensor, eps: float) -> Tensor:
     """
     Normalizes the gramian with respect to the Frobenius norm.
 
@@ -25,7 +25,7 @@ def _normalize(gramian: Tensor, eps: float) -> Tensor:
         return gramian / squared_frobenius_norm
 
 
-def _regularize(gramian: Tensor, eps: float) -> Tensor:
+def regularize(gramian: Tensor, eps: float) -> Tensor:
     """
     Adds a regularization term to the gramian to enforce positive definiteness.
 

src/torchjd/aggregation/cagrad.py

@@ -3,7 +3,7 @@
 import torch
 from torch import Tensor
 
-from ._gramian_utils import _compute_gramian, _normalize
+from ._gramian_utils import compute_gramian, normalize
 from .bases import _WeightedAggregator, _Weighting
 
 
@@ -72,7 +72,7 @@ def __init__(self, c: float, norm_eps: float):
         self.norm_eps = norm_eps
 
     def forward(self, matrix: Tensor) -> Tensor:
-        gramian = _normalize(_compute_gramian(matrix), self.norm_eps)
+        gramian = normalize(compute_gramian(matrix), self.norm_eps)
         U, S, _ = torch.svd(gramian)
 
         reduced_matrix = U @ S.sqrt().diag()

src/torchjd/aggregation/dualproj.py

@@ -3,7 +3,7 @@
 from torch import Tensor
 
 from ._dual_cone_utils import _project_weights
-from ._gramian_utils import _compute_gramian, _normalize, _regularize
+from ._gramian_utils import compute_gramian, normalize, regularize
 from ._pref_vector_utils import _pref_vector_to_str_suffix, _pref_vector_to_weighting
 from .bases import _WeightedAggregator, _Weighting
 from .mean import _MeanWeighting
@@ -100,6 +100,6 @@ def __init__(
 
     def forward(self, matrix: Tensor) -> Tensor:
         u = self.weighting(matrix)
-        G = _regularize(_normalize(_compute_gramian(matrix), self.norm_eps), self.reg_eps)
+        G = regularize(normalize(compute_gramian(matrix), self.norm_eps), self.reg_eps)
         w = _project_weights(u, G, self.solver)
         return w

src/torchjd/aggregation/mgda.py

@@ -1,7 +1,7 @@
 import torch
 from torch import Tensor
 
-from ._gramian_utils import _compute_gramian
+from ._gramian_utils import compute_gramian
 from .bases import _WeightedAggregator, _Weighting
 
 
@@ -57,7 +57,7 @@ def __init__(self, epsilon: float, max_iters: int):
         self.max_iters = max_iters
 
     def _frank_wolfe_solver(self, matrix: Tensor) -> Tensor:
-        gramian = _compute_gramian(matrix)
+        gramian = compute_gramian(matrix)
         device = matrix.device
         dtype = matrix.dtype
 

src/torchjd/aggregation/upgrad.py

@@ -4,7 +4,7 @@
 from torch import Tensor
 
 from ._dual_cone_utils import _project_weights
-from ._gramian_utils import _compute_gramian, _normalize, _regularize
+from ._gramian_utils import compute_gramian, normalize, regularize
 from ._pref_vector_utils import _pref_vector_to_str_suffix, _pref_vector_to_weighting
 from .bases import _WeightedAggregator, _Weighting
 from .mean import _MeanWeighting
@@ -96,6 +96,6 @@ def __init__(
 
     def forward(self, matrix: Tensor) -> Tensor:
         U = torch.diag(self.weighting(matrix))
-        G = _regularize(_normalize(_compute_gramian(matrix), self.norm_eps), self.reg_eps)
+        G = regularize(normalize(compute_gramian(matrix), self.norm_eps), self.reg_eps)
         W = _project_weights(U, G, self.solver)
         return torch.sum(W, dim=0)