Remove _compute_normalized_regularized_gramian and _compute_normalized_gramian as they are now just combinations of _compute_gramian, _regularize and _normalize

Author: PierreQuinton

src/torchjd/aggregation/_gramian_utils.py

@@ -10,11 +10,6 @@ def _compute_gramian(matrix: Tensor) -> Tensor:
     return matrix @ matrix.T
 
 
-def _compute_regularized_normalized_gramian(matrix: Tensor, norm_eps: float, reg_eps: float):
-    normalized_gramian = _compute_normalized_gramian(matrix, norm_eps)
-    return _regularize(normalized_gramian, reg_eps)
-
-
 def _compute_normalized_gramian(matrix: Tensor, eps: float) -> Tensor:
     gramian = _compute_gramian(matrix)
     return _normalize(gramian, eps)

src/torchjd/aggregation/cagrad.py

@@ -3,7 +3,7 @@
 import torch
 from torch import Tensor
 
-from ._gramian_utils import _compute_normalized_gramian
+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 = _compute_normalized_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_regularized_normalized_gramian
+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 = _compute_regularized_normalized_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/upgrad.py

@@ -4,7 +4,7 @@
 from torch import Tensor
 
 from ._dual_cone_utils import _project_weights
-from ._gramian_utils import _compute_regularized_normalized_gramian
+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 = _compute_regularized_normalized_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)