refactor: Improve PSD typing

src/torchjd/_linalg/__init__.py

@@ -1,4 +1,23 @@
-from .gramian import compute_gramian
-from .matrix import Matrix, PSDMatrix
+from ._gramian import compute_gramian
+from ._matrix import (
+    GeneralizedMatrix,
+    Matrix,
+    PSDMatrix,
+    PSDQuadraticForm,
+    is_generalized_matrix,
+    is_matrix,
+    is_psd_matrix,
+    is_psd_quadratic_form,
+)
 
-__all__ = ["compute_gramian", "Matrix", "PSDMatrix"]
+__all__ = [
+    "compute_gramian",
+    "GeneralizedMatrix",
+    "Matrix",
+    "PSDMatrix",
+    "PSDQuadraticForm",
+    "is_generalized_matrix",
+    "is_matrix",
+    "is_psd_matrix",
+    "is_psd_quadratic_form",
+]

src/torchjd/_linalg/gramian.py → src/torchjd/_linalg/_gramian.py

@@ -1,9 +1,11 @@
-from .matrix import Matrix, PSDMatrix
+from ._matrix import Matrix, PSDMatrix, is_psd_matrix
 
 
 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
+    gramian = matrix @ matrix.T
+    assert is_psd_matrix(gramian)
+    return gramian

src/torchjd/_linalg/_matrix.py

@@ -0,0 +1,42 @@
+from typing import TypeGuard
+
+from torch import Tensor
+
+
+class GeneralizedMatrix(Tensor):
+    pass
+
+
+class Matrix(GeneralizedMatrix):
+    pass
+
+
+class PSDQuadraticForm(Tensor):
+    pass
+
+
+class PSDMatrix(PSDQuadraticForm, Matrix):
+    pass
+
+
+def is_generalized_matrix(t: Tensor) -> TypeGuard[GeneralizedMatrix]:
+    return t.ndim >= 1
+
+
+def is_matrix(t: Tensor) -> TypeGuard[Matrix]:
+    return t.ndim == 2
+
+
+def is_psd_quadratic_form(t: Tensor) -> TypeGuard[PSDQuadraticForm]:
+    half_dim = t.ndim // 2
+    return not t.ndim % 2 != 0 and t.shape[:half_dim] == t.shape[: half_dim - 1 : -1]
+    # We do not check that t is PSD as it is expensive, but this must be checked in the tests of
+    # every function that use this TypeGuard.
+    # TODO: Say with what assert we check that
+
+
+def is_psd_matrix(t: Tensor) -> TypeGuard[PSDMatrix]:
+    return t.ndim == 2 and t.shape[0] == t.shape[1]
+    # We do not check that t is PSD as it is expensive, but this must be checked in the tests of
+    # every function that use this TypeGuard.
+    # TODO: Say with what assert we check that

src/torchjd/aggregation/_aggregator_bases.py

@@ -2,7 +2,7 @@
 
 from torch import Tensor, nn
 
-from torchjd._linalg import Matrix, PSDMatrix, compute_gramian
+from torchjd._linalg import Matrix, PSDMatrix, compute_gramian, is_matrix
 
 from ._weighting_bases import Weighting
 
@@ -18,20 +18,21 @@ def __init__(self):
 
     @staticmethod
     def _check_is_matrix(matrix: Tensor) -> None:
-        if len(matrix.shape) != 2:
+        if not is_matrix(matrix):
             raise ValueError(
                 "Parameter `matrix` should be a tensor of dimension 2. Found `matrix.shape = "
                 f"{matrix.shape}`."
             )
 
     @abstractmethod
-    def forward(self, matrix: Tensor) -> Tensor:
+    def forward(self, matrix: Matrix) -> Tensor:
         """Computes the aggregation from the input matrix."""
 
-    # Override to make type hints and documentation more specific
+    # Override to make type hints and documentation more specific, note that `Matrix` type isn't
+    # public
     def __call__(self, matrix: Tensor) -> Tensor:
         """Computes the aggregation from the input matrix and applies all registered hooks."""
-
+        Aggregator._check_is_matrix(matrix)
         return super().__call__(matrix)
 
     def __repr__(self) -> str:
@@ -54,7 +55,7 @@ def __init__(self, weighting: Weighting[Matrix]):
         self.weighting = weighting
 
     @staticmethod
-    def combine(matrix: Tensor, weights: Tensor) -> Tensor:
+    def combine(matrix: Matrix, weights: Tensor) -> Tensor:
         """
         Aggregates a matrix by making a linear combination of its rows, using the provided vector of
         weights.
@@ -63,8 +64,7 @@ def combine(matrix: Tensor, weights: Tensor) -> Tensor:
         vector = weights @ matrix
         return vector
 
-    def forward(self, matrix: Tensor) -> Tensor:
-        self._check_is_matrix(matrix)
+    def forward(self, matrix: Matrix) -> Tensor:
         weights = self.weighting(matrix)
         vector = self.combine(matrix, weights)
         return vector

src/torchjd/aggregation/_config.py

@@ -28,6 +28,8 @@
 import torch
 from torch import Tensor
 
+from torchjd._linalg import Matrix
+
 from ._aggregator_bases import Aggregator
 from ._sum import SumWeighting
 from ._utils.non_differentiable import raise_non_differentiable_error
@@ -56,7 +58,7 @@ def __init__(self, pref_vector: Tensor | None = None):
         # This prevents computing gradients that can be very wrong.
         self.register_full_backward_pre_hook(raise_non_differentiable_error)
 
-    def forward(self, matrix: Tensor) -> Tensor:
+    def forward(self, matrix: Matrix) -> Tensor:
         weights = self.weighting(matrix)
         units = torch.nan_to_num((matrix / (matrix.norm(dim=1)).unsqueeze(1)), 0.0)
         best_direction = torch.linalg.pinv(units) @ weights

src/torchjd/aggregation/_flattening.py

@@ -2,7 +2,7 @@
 
 from torch import Tensor
 
-from torchjd._linalg.matrix import PSDMatrix
+from torchjd._linalg import PSDMatrix, PSDQuadraticForm, is_psd_matrix
 from torchjd.aggregation._weighting_bases import GeneralizedWeighting, Weighting
 from torchjd.autogram._gramian_utils import reshape_gramian
 
@@ -26,11 +26,12 @@ def __init__(self, weighting: Weighting[PSDMatrix]):
         super().__init__()
         self.weighting = weighting
 
-    def forward(self, generalized_gramian: Tensor) -> Tensor:
+    def forward(self, generalized_gramian: PSDQuadraticForm) -> Tensor:
         k = generalized_gramian.ndim // 2
         shape = generalized_gramian.shape[:k]
         m = prod(shape)
         square_gramian = reshape_gramian(generalized_gramian, [m])
+        assert is_psd_matrix(square_gramian)
         weights_vector = self.weighting(square_gramian)
         weights = weights_vector.reshape(shape)
         return weights

src/torchjd/aggregation/_graddrop.py

@@ -3,6 +3,8 @@
 import torch
 from torch import Tensor
 
+from torchjd._linalg import Matrix
+
 from ._aggregator_bases import Aggregator
 from ._utils.non_differentiable import raise_non_differentiable_error
 
@@ -38,8 +40,7 @@ def __init__(self, f: Callable = _identity, leak: Tensor | None = None):
         # This prevents computing gradients that can be very wrong.
         self.register_full_backward_pre_hook(raise_non_differentiable_error)
 
-    def forward(self, matrix: Tensor) -> Tensor:
-        self._check_is_matrix(matrix)
+    def forward(self, matrix: Matrix) -> Tensor:
         self._check_matrix_has_enough_rows(matrix)
 
         if matrix.shape[0] == 0 or matrix.shape[1] == 0:

src/torchjd/aggregation/_pcgrad.py

@@ -1,3 +1,5 @@
+from typing import cast
+
 import torch
 from torch import Tensor
 
@@ -32,7 +34,7 @@ def forward(self, gramian: PSDMatrix) -> Tensor:
         device = gramian.device
         dtype = gramian.dtype
         cpu = torch.device("cpu")
-        gramian = gramian.to(device=cpu)
+        gramian = cast(PSDMatrix, gramian.to(device=cpu))
 
         dimension = gramian.shape[0]
         weights = torch.zeros(dimension, device=cpu, dtype=dtype)

src/torchjd/aggregation/_trimmed_mean.py

@@ -25,7 +25,6 @@ def __init__(self, trim_number: int):
         self.trim_number = trim_number
 
     def forward(self, matrix: Tensor) -> Tensor:
-        self._check_is_matrix(matrix)
         self._check_matrix_has_enough_rows(matrix)
 
         n_rows = matrix.shape[0]

src/torchjd/aggregation/_utils/gramian.py

@@ -1,6 +1,6 @@
 import torch
 
-from torchjd._linalg.matrix import PSDMatrix
+from torchjd._linalg import PSDMatrix, is_psd_matrix
 
 
 def normalize(gramian: PSDMatrix, eps: float) -> PSDMatrix:
@@ -13,9 +13,11 @@ def normalize(gramian: PSDMatrix, eps: float) -> PSDMatrix:
     """
     squared_frobenius_norm = gramian.diagonal().sum()
     if squared_frobenius_norm < eps:
-        return torch.zeros_like(gramian)
+        output = torch.zeros_like(gramian)
     else:
-        return gramian / squared_frobenius_norm
+        output = gramian / squared_frobenius_norm
+    assert is_psd_matrix(output)
+    return output
 
 
 def regularize(gramian: PSDMatrix, eps: float) -> PSDMatrix:
@@ -30,4 +32,6 @@ def regularize(gramian: PSDMatrix, eps: float) -> PSDMatrix:
     regularization_matrix = eps * torch.eye(
         gramian.shape[0], dtype=gramian.dtype, device=gramian.device
     )
-    return gramian + regularization_matrix
+    output = gramian + regularization_matrix
+    assert is_psd_matrix(output)
+    return output

src/torchjd/aggregation/_utils/pref_vector.py

@@ -1,6 +1,6 @@
 from torch import Tensor
 
-from torchjd._linalg.matrix import Matrix
+from torchjd._linalg import Matrix
 from torchjd.aggregation._constant import ConstantWeighting
 from torchjd.aggregation._weighting_bases import Weighting
 

src/torchjd/aggregation/_weighting_bases.py

@@ -6,6 +6,8 @@
 
 from torch import Tensor, nn
 
+from torchjd._linalg import PSDQuadraticForm, is_psd_quadratic_form
+
 _T = TypeVar("_T", contravariant=True)
 _FnInputT = TypeVar("_FnInputT")
 _FnOutputT = TypeVar("_FnOutputT")
@@ -64,7 +66,7 @@ def __init__(self):
         super().__init__()
 
     @abstractmethod
-    def forward(self, generalized_gramian: Tensor) -> Tensor:
+    def forward(self, generalized_gramian: PSDQuadraticForm) -> Tensor:
         """Computes the vector of weights from the input generalized Gramian."""
 
     # Override to make type hints and documentation more specific
@@ -74,4 +76,5 @@ def __call__(self, generalized_gramian: Tensor) -> Tensor:
         hooks.
         """
 
+        assert is_psd_quadratic_form(generalized_gramian)
         return super().__call__(generalized_gramian)

src/torchjd/autogram/_engine.py

@@ -4,7 +4,7 @@
 from torch import Tensor, nn, vmap
 from torch.autograd.graph import get_gradient_edge
 
-from torchjd._linalg.matrix import PSDMatrix
+from torchjd._linalg import PSDMatrix
 
 from ._edge_registry import EdgeRegistry
 from ._gramian_accumulator import GramianAccumulator

src/torchjd/autogram/_gramian_accumulator.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from torchjd._linalg.matrix import PSDMatrix
+from torchjd._linalg import PSDMatrix
 
 
 class GramianAccumulator:

src/torchjd/autogram/_gramian_computer.py

@@ -4,8 +4,7 @@
 from torch import Tensor
 from torch.utils._pytree import PyTree
 
-from torchjd._linalg import compute_gramian
-from torchjd._linalg.matrix import PSDMatrix
+from torchjd._linalg import Matrix, PSDMatrix, compute_gramian, is_matrix
 from torchjd.autogram._jacobian_computer import JacobianComputer
 
 
@@ -23,12 +22,12 @@ def __call__(
     def track_forward_call(self) -> None:
         """Track that the module's forward was called. Necessary in some implementations."""
 
-    def reset(self):
+    def reset(self) -> None:
         """Reset state if any. Necessary in some implementations."""
 
 
 class JacobianBasedGramianComputer(GramianComputer, ABC):
-    def __init__(self, jacobian_computer):
+    def __init__(self, jacobian_computer: JacobianComputer):
         self.jacobian_computer = jacobian_computer
 
 
@@ -41,7 +40,7 @@ class JacobianBasedGramianComputerWithCrossTerms(JacobianBasedGramianComputer):
     def __init__(self, jacobian_computer: JacobianComputer):
         super().__init__(jacobian_computer)
         self.remaining_counter = 0
-        self.summed_jacobian: Optional[Tensor] = None
+        self.summed_jacobian: Optional[Matrix] = None
 
     def reset(self) -> None:
         self.remaining_counter = 0
@@ -64,7 +63,9 @@ def __call__(
         if self.summed_jacobian is None:
             self.summed_jacobian = jacobian_matrix
         else:
-            self.summed_jacobian += jacobian_matrix
+            jacobians_sum = self.summed_jacobian + jacobian_matrix
+            assert is_matrix(jacobians_sum)
+            self.summed_jacobian = jacobians_sum
 
         self.remaining_counter -= 1