refactor(linalg): Create _linalg package (#520)

* Create `_linalg` package, factorize computation of Gramian and put resulting `compute_gramian` in `_linalg`.

* Move Matrix and PSDMatrix to `_linalg/matrix.py`

* Properly type gramians as `PSDMatrix`.

* Add assert_psd_matrix

* test that compute_gramian returns a psd_matrix

Author: PierreQuinton

src/torchjd/_linalg/__init__.py

@@ -0,0 +1,4 @@
+from .gramian import compute_gramian
+from .matrix import Matrix, PSDMatrix
+
+__all__ = ["compute_gramian", "Matrix", "PSDMatrix"]

src/torchjd/_linalg/gramian.py

@@ -0,0 +1,9 @@
+from .matrix import Matrix, PSDMatrix
+
+
+def compute_gramian(matrix: Matrix) -> PSDMatrix:
+    """
+    Computes the `Gramian matrix <https://en.wikipedia.org/wiki/Gram_matrix>`_ of a given matrix.
+    """
+
+    return matrix @ matrix.T

src/torchjd/_linalg/matrix.py

@@ -0,0 +1,6 @@
+from typing import Annotated
+
+from torch import Tensor
+
+Matrix = Annotated[Tensor, "ndim=2"]
+PSDMatrix = Annotated[Matrix, "Positive semi-definite"]

src/torchjd/aggregation/_aggregator_bases.py

@@ -2,8 +2,9 @@
 
 from torch import Tensor, nn
 
-from ._utils.gramian import compute_gramian
-from ._weighting_bases import Matrix, PSDMatrix, Weighting
+from torchjd._linalg import Matrix, PSDMatrix, compute_gramian
+
+from ._weighting_bases import Weighting
 
 
 class Aggregator(nn.Module, ABC):

src/torchjd/aggregation/_aligned_mtl.py

@@ -28,10 +28,12 @@
 import torch
 from torch import Tensor
 
+from torchjd._linalg import PSDMatrix
+
 from ._aggregator_bases import GramianWeightedAggregator
 from ._mean import MeanWeighting
 from ._utils.pref_vector import pref_vector_to_str_suffix, pref_vector_to_weighting
-from ._weighting_bases import PSDMatrix, Weighting
+from ._weighting_bases import Weighting
 
 
 class AlignedMTL(GramianWeightedAggregator):
@@ -73,7 +75,7 @@ def __init__(self, pref_vector: Tensor | None = None):
         self._pref_vector = pref_vector
         self.weighting = pref_vector_to_weighting(pref_vector, default=MeanWeighting())
 
-    def forward(self, gramian: Tensor) -> Tensor:
+    def forward(self, gramian: PSDMatrix) -> Tensor:
         w = self.weighting(gramian)
         B = self._compute_balance_transformation(gramian)
         alpha = B @ w

src/torchjd/aggregation/_cagrad.py

@@ -1,7 +1,9 @@
 from typing import cast
 
+from torchjd._linalg import PSDMatrix
+
 from ._utils.check_dependencies import check_dependencies_are_installed
-from ._weighting_bases import PSDMatrix, Weighting
+from ._weighting_bases import Weighting
 
 check_dependencies_are_installed(["cvxpy", "clarabel"])
 
@@ -73,7 +75,7 @@ def __init__(self, c: float, norm_eps: float = 0.0001):
         self.c = c
         self.norm_eps = norm_eps
 
-    def forward(self, gramian: Tensor) -> Tensor:
+    def forward(self, gramian: PSDMatrix) -> Tensor:
         U, S, _ = torch.svd(normalize(gramian, self.norm_eps))
 
         reduced_matrix = U @ S.sqrt().diag()

src/torchjd/aggregation/_constant.py

@@ -1,8 +1,10 @@
 from torch import Tensor
 
+from torchjd._linalg import Matrix
+
 from ._aggregator_bases import WeightedAggregator
 from ._utils.str import vector_to_str
-from ._weighting_bases import Matrix, Weighting
+from ._weighting_bases import Weighting
 
 
 class Constant(WeightedAggregator):

src/torchjd/aggregation/_dualproj.py

@@ -2,13 +2,15 @@
 
 from torch import Tensor
 
+from torchjd._linalg import PSDMatrix
+
 from ._aggregator_bases import GramianWeightedAggregator
 from ._mean import MeanWeighting
 from ._utils.dual_cone import project_weights
 from ._utils.gramian import normalize, regularize
 from ._utils.non_differentiable import raise_non_differentiable_error
 from ._utils.pref_vector import pref_vector_to_str_suffix, pref_vector_to_weighting
-from ._weighting_bases import PSDMatrix, Weighting
+from ._weighting_bases import Weighting
 
 
 class DualProj(GramianWeightedAggregator):
@@ -86,7 +88,7 @@ def __init__(
         self.reg_eps = reg_eps
         self.solver = solver
 
-    def forward(self, gramian: Tensor) -> Tensor:
+    def forward(self, gramian: PSDMatrix) -> Tensor:
         u = self.weighting(gramian)
         G = regularize(normalize(gramian, self.norm_eps), self.reg_eps)
         w = project_weights(u, G, self.solver)

src/torchjd/aggregation/_flattening.py

@@ -2,7 +2,8 @@
 
 from torch import Tensor
 
-from torchjd.aggregation._weighting_bases import GeneralizedWeighting, PSDMatrix, Weighting
+from torchjd._linalg.matrix import PSDMatrix
+from torchjd.aggregation._weighting_bases import GeneralizedWeighting, Weighting
 from torchjd.autogram._gramian_utils import reshape_gramian
 
 

src/torchjd/aggregation/_imtl_g.py

@@ -1,9 +1,11 @@
 import torch
 from torch import Tensor
 
+from torchjd._linalg import PSDMatrix
+
 from ._aggregator_bases import GramianWeightedAggregator
 from ._utils.non_differentiable import raise_non_differentiable_error
-from ._weighting_bases import PSDMatrix, Weighting
+from ._weighting_bases import Weighting
 
 
 class IMTLG(GramianWeightedAggregator):
@@ -27,7 +29,7 @@ class IMTLGWeighting(Weighting[PSDMatrix]):
     :class:`~torchjd.aggregation.IMTLG`.
     """
 
-    def forward(self, gramian: Tensor) -> Tensor:
+    def forward(self, gramian: PSDMatrix) -> Tensor:
         d = torch.sqrt(torch.diagonal(gramian))
         v = torch.linalg.pinv(gramian) @ d
         v_sum = v.sum()