Apply suggestions from code review

Co-authored-by: Valérian Rey <31951177+ValerianRey@users.noreply.github.com>

Author: PierreQuinton

src/torchjd/_linalg/_matrix.py

@@ -4,19 +4,23 @@
 
 
 class GeneralizedMatrix(Tensor):
-    pass
+    """Tensor with a least 1 dimension."""
 
 
 class Matrix(GeneralizedMatrix):
-    pass
+    """Tensor with exactly 2 dimensions."""
 
 
 class PSDGeneralizedMatrix(Tensor):
-    pass
+    """
+    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):
-    pass
+    """Positive semi-definite matrix."""
 
 
 def is_generalized_matrix(t: Tensor) -> TypeGuard[GeneralizedMatrix]:
@@ -31,10 +35,10 @@ def is_psd_generalized_matrix(t: Tensor) -> TypeGuard[PSDGeneralizedMatrix]:
     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 use this TypeGuard by using `assert_psd_generalized_matrix`.
+    # every function that uses this TypeGuard by using `assert_psd_generalized_matrix`.
 
 
 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, by using `assert_psd_matrix`.
+    # every function that uses this TypeGuard, by using `assert_psd_matrix`.

src/torchjd/autogram/_gramian_utils.py

@@ -16,11 +16,11 @@ def reshape(gramian: PSDGeneralizedMatrix, half_shape: list[int]) -> PSDGenerali
         `shape + shape[::-1]`.
     """
 
-    # Example 1: `gramian` of shape [4, 3, 2, 2, 3, 4] and `shape` of [8, 3]:
+    # Example 1: `gramian` of shape [4, 3, 2, 2, 3, 4] and `half_shape` of [8, 3]:
     # [4, 3, 2, 2, 3, 4] -(movedim)-> [4, 3, 2, 4, 3, 2] -(reshape)-> [8, 3, 8, 3] -(movedim)->
     # [8, 3, 3, 8]
     #
-    # Example 2: `gramian` of shape [24, 24] and `shape` of [4, 3, 2]:
+    # Example 2: `gramian` of shape [24, 24] and `half_shape` of [4, 3, 2]:
     # [24, 24] -(movedim)-> [24, 24] -(reshape)-> [4, 3, 2, 4, 3, 2] -(movedim)-> [4, 3, 2, 2, 3, 4]
 
     reshaped_gramian = _revert_last_dims(
@@ -31,8 +31,8 @@ def reshape(gramian: PSDGeneralizedMatrix, half_shape: list[int]) -> PSDGenerali
 
 def flatten(gramian: PSDGeneralizedMatrix) -> PSDMatrix:
     k = gramian.ndim // 2
-    shape = gramian.shape[:k]
-    m = prod(shape)
+    half_shape = gramian.shape[:k]
+    m = prod(half_shape)
     square_gramian = reshape(gramian, [m])
     return cast(PSDMatrix, square_gramian)