Rename PSDGeneralizedMatrix to PSDTensor

Author: ValerianRey

src/torchjd/_linalg/__init__.py

@@ -1,21 +1,14 @@
 from ._gramian import compute_gramian, normalize, regularize
-from ._matrix import (
-    Matrix,
-    PSDGeneralizedMatrix,
-    PSDMatrix,
-    is_matrix,
-    is_psd_generalized_matrix,
-    is_psd_matrix,
-)
+from ._matrix import Matrix, PSDMatrix, PSDTensor, is_matrix, is_psd_matrix, is_psd_tensor
 
 __all__ = [
     "compute_gramian",
     "normalize",
     "regularize",
     "Matrix",
     "PSDMatrix",
-    "PSDGeneralizedMatrix",
+    "PSDTensor",
     "is_matrix",
     "is_psd_matrix",
-    "is_psd_generalized_matrix",
+    "is_psd_tensor",
 ]

src/torchjd/_linalg/_gramian.py

@@ -3,7 +3,7 @@
 import torch
 from torch import Tensor
 
-from ._matrix import PSDGeneralizedMatrix, PSDMatrix
+from ._matrix import PSDMatrix, PSDTensor
 
 
 @overload
@@ -17,11 +17,11 @@ def compute_gramian(matrix: Tensor, contracted_dims: Literal[-1]) -> PSDMatrix:
 
 
 @overload
-def compute_gramian(matrix: Tensor, contracted_dims: int) -> PSDGeneralizedMatrix:
+def compute_gramian(matrix: Tensor, contracted_dims: int) -> PSDTensor:
     pass
 
 
-def compute_gramian(matrix: Tensor, contracted_dims: int = -1) -> PSDGeneralizedMatrix:
+def compute_gramian(matrix: Tensor, contracted_dims: int = -1) -> PSDTensor:
     """
     Computes the `Gramian matrix <https://en.wikipedia.org/wiki/Gram_matrix>`_ of the input.
 
@@ -35,7 +35,7 @@ def compute_gramian(matrix: Tensor, contracted_dims: int = -1) -> PSDGeneralized
     indices_dest = list(range(matrix.ndim - 1, contracted_dims - 1, -1))
     transposed_matrix = matrix.movedim(indices_source, indices_dest)
     gramian = torch.tensordot(matrix, transposed_matrix, dims=contracted_dims)
-    return cast(PSDGeneralizedMatrix, gramian)
+    return cast(PSDTensor, gramian)
 
 
 def normalize(gramian: PSDMatrix, eps: float) -> PSDMatrix:

src/torchjd/_linalg/_matrix.py

@@ -7,30 +7,30 @@ class Matrix(Tensor):
     """Tensor with exactly 2 dimensions."""
 
 
-class PSDGeneralizedMatrix(Tensor):
+class PSDTensor(Tensor):
     """
     Tensor representing a quadratic form. The first half of its dimensions matches the reversed
     second half of its dimensions (e.g. shape=[4, 3, 3, 4]), and its reshaping into a matrix should
     be positive semi-definite.
     """
 
 
-class PSDMatrix(PSDGeneralizedMatrix, Matrix):
+class PSDMatrix(PSDTensor, Matrix):
     """Positive semi-definite matrix."""
 
 
 def is_matrix(t: Tensor) -> TypeGuard[Matrix]:
     return t.ndim == 2
 
 
-def is_psd_generalized_matrix(t: Tensor) -> TypeGuard[PSDGeneralizedMatrix]:
+def is_psd_tensor(t: Tensor) -> TypeGuard[PSDTensor]:
     half_dim = t.ndim // 2
     return 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 uses this TypeGuard by using `assert_psd_generalized_matrix`.
+    # every function that uses this TypeGuard by using `assert_is_psd_tensor`.
 
 
 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 uses this TypeGuard, by using `assert_psd_matrix`.
+    # every function that uses this TypeGuard, by using `assert_is_psd_matrix`.

src/torchjd/aggregation/_flattening.py

@@ -1,6 +1,6 @@
 from torch import Tensor
 
-from torchjd._linalg import PSDGeneralizedMatrix
+from torchjd._linalg import PSDTensor
 from torchjd.aggregation._weighting_bases import GeneralizedWeighting, Weighting
 from torchjd.autogram._gramian_utils import flatten
 
@@ -24,7 +24,7 @@ def __init__(self, weighting: Weighting):
         super().__init__()
         self.weighting = weighting
 
-    def forward(self, generalized_gramian: PSDGeneralizedMatrix) -> Tensor:
+    def forward(self, generalized_gramian: PSDTensor) -> Tensor:
         k = generalized_gramian.ndim // 2
         shape = generalized_gramian.shape[:k]
         square_gramian = flatten(generalized_gramian)

src/torchjd/aggregation/_weighting_bases.py

@@ -6,7 +6,7 @@
 
 from torch import Tensor, nn
 
-from torchjd._linalg import PSDGeneralizedMatrix, is_psd_generalized_matrix
+from torchjd._linalg import PSDTensor, is_psd_tensor
 
 _T = TypeVar("_T", contravariant=True)
 _FnInputT = TypeVar("_FnInputT")
@@ -65,7 +65,7 @@ def __init__(self):
         super().__init__()
 
     @abstractmethod
-    def forward(self, generalized_gramian: PSDGeneralizedMatrix) -> Tensor:
+    def forward(self, generalized_gramian: PSDTensor) -> Tensor:
         """Computes the vector of weights from the input generalized Gramian."""
 
     def __call__(self, generalized_gramian: Tensor) -> Tensor:
@@ -74,5 +74,5 @@ def __call__(self, generalized_gramian: Tensor) -> Tensor:
         hooks.
         """
 
-        assert is_psd_generalized_matrix(generalized_gramian)
+        assert is_psd_tensor(generalized_gramian)
         return super().__call__(generalized_gramian)

src/torchjd/autogram/_gramian_utils.py

@@ -3,10 +3,10 @@
 
 from torch import Tensor
 
-from torchjd._linalg import PSDGeneralizedMatrix, PSDMatrix
+from torchjd._linalg import PSDMatrix, PSDTensor
 
 
-def flatten(gramian: PSDGeneralizedMatrix) -> PSDMatrix:
+def flatten(gramian: PSDTensor) -> PSDMatrix:
     """
     Flattens a generalized Gramian into a square matrix. The first half of the dimensions are
     flattened into the first dimension, and the second half are flattened into the second.
@@ -24,7 +24,7 @@ def flatten(gramian: PSDGeneralizedMatrix) -> PSDMatrix:
     return cast(PSDMatrix, square_gramian)
 
 
-def reshape(gramian: PSDGeneralizedMatrix, half_shape: list[int]) -> PSDGeneralizedMatrix:
+def reshape(gramian: PSDTensor, half_shape: list[int]) -> PSDTensor:
     """
     Reshapes a Gramian to a provided shape. The reshape of the first half of the target dimensions
     must be done from the left, while the reshape of the second half must be done from the right.
@@ -42,7 +42,7 @@ def reshape(gramian: PSDGeneralizedMatrix, half_shape: list[int]) -> PSDGenerali
     # [24, 24] -(movedim)-> [24, 24] -(reshape)-> [4, 3, 2, 4, 3, 2] -(movedim)-> [4, 3, 2, 2, 3, 4]
 
     result = _revert_last_dims(_revert_last_dims(gramian).reshape(half_shape + half_shape))
-    return cast(PSDGeneralizedMatrix, result)
+    return cast(PSDTensor, result)
 
 
 def _revert_last_dims(t: Tensor) -> Tensor:
@@ -53,9 +53,7 @@ def _revert_last_dims(t: Tensor) -> Tensor:
     return t.movedim(last_dims, last_dims[::-1])
 
 
-def movedim(
-    gramian: PSDGeneralizedMatrix, half_source: list[int], half_destination: list[int]
-) -> PSDGeneralizedMatrix:
+def movedim(gramian: PSDTensor, half_source: list[int], half_destination: list[int]) -> PSDTensor:
     """
     Moves the dimensions of a Gramian from some source dimensions to destination dimensions. This
     must be done simultaneously on the first half of the dimensions and on the second half of the
@@ -86,4 +84,4 @@ def movedim(
     source = half_source_ + [last_dim - i for i in half_source_]
     destination = half_destination_ + [last_dim - i for i in half_destination_]
     moved_gramian = gramian.movedim(source, destination)
-    return cast(PSDGeneralizedMatrix, moved_gramian)
+    return cast(PSDTensor, moved_gramian)

tests/unit/autogram/test_gramian_utils.py

@@ -1,6 +1,6 @@
 from pytest import mark
 from torch.testing import assert_close
-from utils.asserts import assert_is_psd_generalized_matrix, assert_is_psd_matrix
+from utils.asserts import assert_is_psd_matrix, assert_is_psd_tensor
 from utils.tensors import randn_
 
 from torchjd._linalg import compute_gramian, is_psd_matrix
@@ -59,7 +59,7 @@ def test_reshape_yields_psd(original_shape: list[int], target_shape: list[int]):
     matrix = randn_(original_shape + [2])
     gramian = compute_gramian(matrix, 1)
     reshaped_gramian = reshape(gramian, target_shape)
-    assert_is_psd_generalized_matrix(reshaped_gramian, atol=1e-04, rtol=0.0)
+    assert_is_psd_tensor(reshaped_gramian, atol=1e-04, rtol=0.0)
 
 
 @mark.parametrize(
@@ -150,4 +150,4 @@ def test_movedim_yields_psd(shape: list[int], source: list[int], destination: li
     matrix = randn_(shape + [2])
     gramian = compute_gramian(matrix, 1)
     moveddim_gramian = movedim(gramian, source, destination)
-    assert_is_psd_generalized_matrix(moveddim_gramian)
+    assert_is_psd_tensor(moveddim_gramian)

tests/utils/asserts.py

@@ -2,7 +2,7 @@
 from torch import Tensor
 from torch.testing import assert_close
 
-from torchjd._linalg import is_psd_generalized_matrix, is_psd_matrix
+from torchjd._linalg import is_psd_matrix, is_psd_tensor
 from torchjd.autogram._gramian_utils import flatten
 from torchjd.autojac._accumulation import is_tensor_with_jac
 
@@ -44,7 +44,7 @@ def assert_is_psd_matrix(matrix: Tensor, **kwargs) -> None:
     assert_close(eig_vals, expected_eig_vals, **kwargs)
 
 
-def assert_is_psd_generalized_matrix(t: Tensor, **kwargs) -> None:
-    assert is_psd_generalized_matrix(t)
+def assert_is_psd_tensor(t: Tensor, **kwargs) -> None:
+    assert is_psd_tensor(t)
     matrix = flatten(t)
     assert_is_psd_matrix(matrix, **kwargs)

tests/utils/forward_backwards.py

@@ -8,7 +8,7 @@
 from utils.architectures import get_in_out_shapes
 from utils.contexts import fork_rng
 
-from torchjd._linalg import PSDGeneralizedMatrix
+from torchjd._linalg import PSDTensor
 from torchjd.aggregation import Aggregator, Weighting
 from torchjd.autogram import Engine
 from torchjd.autojac import backward
@@ -117,7 +117,7 @@ def reshape_raw_losses(raw_losses: Tensor) -> Tensor:
 
 def compute_gramian_with_autograd(
     output: Tensor, params: list[nn.Parameter], retain_graph: bool = False
-) -> PSDGeneralizedMatrix:
+) -> PSDTensor:
     """
     Computes the Gramian of the Jacobian of the outputs with respect to the params using vmapped
     calls to the autograd engine.