Add overload for compute_gramian when t is matrix and contracted_dims is 1

Author: ValerianRey

src/torchjd/_linalg/_gramian.py

@@ -3,7 +3,7 @@
 import torch
 from torch import Tensor
 
-from ._matrix import PSDMatrix, PSDTensor
+from ._matrix import Matrix, PSDMatrix, PSDTensor
 
 
 @overload
@@ -16,6 +16,11 @@ def compute_gramian(t: Tensor, contracted_dims: Literal[-1]) -> PSDMatrix:
     pass
 
 
+@overload
+def compute_gramian(t: Matrix, contracted_dims: Literal[1]) -> PSDMatrix:
+    pass
+
+
 def compute_gramian(t: Tensor, contracted_dims: int = -1) -> PSDTensor:
     """
     Computes the `Gramian matrix <https://en.wikipedia.org/wiki/Gram_matrix>`_ of the input.