Remove GeneralizedMatrix

Author: ValerianRey

src/torchjd/_linalg/__init__.py

@@ -1,10 +1,8 @@
 from ._gramian import compute_gramian, normalize, regularize
 from ._matrix import (
-    GeneralizedMatrix,
     Matrix,
     PSDGeneralizedMatrix,
     PSDMatrix,
-    is_generalized_matrix,
     is_matrix,
     is_psd_generalized_matrix,
     is_psd_matrix,
@@ -14,11 +12,9 @@
     "compute_gramian",
     "normalize",
     "regularize",
-    "GeneralizedMatrix",
     "Matrix",
     "PSDMatrix",
     "PSDGeneralizedMatrix",
-    "is_generalized_matrix",
     "is_matrix",
     "is_psd_matrix",
     "is_psd_generalized_matrix",

src/torchjd/_linalg/_gramian.py

@@ -1,26 +1,27 @@
 from typing import Literal, cast, overload
 
 import torch
+from torch import Tensor
 
-from ._matrix import GeneralizedMatrix, PSDGeneralizedMatrix, PSDMatrix
+from ._matrix import PSDGeneralizedMatrix, PSDMatrix
 
 
 @overload
-def compute_gramian(matrix: GeneralizedMatrix) -> PSDMatrix:
+def compute_gramian(matrix: Tensor) -> PSDMatrix:
     pass
 
 
 @overload
-def compute_gramian(matrix: GeneralizedMatrix, contracted_dims: Literal[-1]) -> PSDMatrix:
+def compute_gramian(matrix: Tensor, contracted_dims: Literal[-1]) -> PSDMatrix:
     pass
 
 
 @overload
-def compute_gramian(matrix: GeneralizedMatrix, contracted_dims: int) -> PSDGeneralizedMatrix:
+def compute_gramian(matrix: Tensor, contracted_dims: int) -> PSDGeneralizedMatrix:
     pass
 
 
-def compute_gramian(matrix: GeneralizedMatrix, contracted_dims: int = -1) -> PSDGeneralizedMatrix:
+def compute_gramian(matrix: Tensor, contracted_dims: int = -1) -> PSDGeneralizedMatrix:
     """
     Computes the `Gramian matrix <https://en.wikipedia.org/wiki/Gram_matrix>`_ of the input.
 

src/torchjd/_linalg/_matrix.py

@@ -3,11 +3,7 @@
 from torch import Tensor
 
 
-class GeneralizedMatrix(Tensor):
-    """Tensor with a least 1 dimension."""
-
-
-class Matrix(GeneralizedMatrix):
+class Matrix(Tensor):
     """Tensor with exactly 2 dimensions."""
 
 
@@ -23,10 +19,6 @@ class PSDMatrix(PSDGeneralizedMatrix, Matrix):
     """Positive semi-definite matrix."""
 
 
-def is_generalized_matrix(t: Tensor) -> TypeGuard[GeneralizedMatrix]:
-    return t.ndim >= 1
-
-
 def is_matrix(t: Tensor) -> TypeGuard[Matrix]:
     return t.ndim == 2
 

tests/unit/linalg/test_gramian.py

@@ -2,7 +2,7 @@
 from utils.asserts import assert_is_psd_matrix
 from utils.tensors import randn_
 
-from torchjd._linalg import compute_gramian, is_generalized_matrix, is_matrix, normalize, regularize
+from torchjd._linalg import compute_gramian, is_matrix, normalize, regularize
 
 
 @mark.parametrize(
@@ -21,7 +21,6 @@
 )
 def test_gramian_is_psd(shape: list[int]):
     matrix = randn_(shape)
-    assert is_generalized_matrix(matrix)
     gramian = compute_gramian(matrix)
     assert_is_psd_matrix(gramian)