refactor(linalg): Add missing overload to compute_gramian (#540)

Author: ValerianRey

src/torchjd/_linalg/_gramian.py

@@ -21,6 +21,11 @@ def compute_gramian(t: Matrix, contracted_dims: Literal[1]) -> PSDMatrix:
     pass
 
 
+@overload
+def compute_gramian(t: Tensor, contracted_dims: int) -> PSDTensor:
+    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.