Improve typing of forward_backward.compute_gramian

Author: PierreQuinton

tests/unit/autogram/test_gramian_utils.py

@@ -4,7 +4,7 @@
 from utils.forward_backwards import compute_gramian
 from utils.tensors import randn_
 
-from torchjd._linalg import is_psd_generalized_matrix, is_psd_matrix
+from torchjd._linalg import is_psd_matrix
 from torchjd.autogram._gramian_utils import flatten, movedim, reshape
 
 
@@ -36,7 +36,6 @@ def test_reshape_equivarience(original_shape: list[int], target_shape: list[int]
     original_gramian = compute_gramian(original_matrix)
     target_gramian = compute_gramian(target_matrix)
 
-    assert is_psd_generalized_matrix(original_gramian)
     reshaped_gramian = reshape(original_gramian, target_shape)
 
     assert_close(reshaped_gramian, target_gramian)
@@ -60,7 +59,6 @@ def test_reshape_equivarience(original_shape: list[int], target_shape: list[int]
 def test_reshape_yields_psd(original_shape: list[int], target_shape: list[int]):
     matrix = randn_(original_shape + [2])
     gramian = compute_gramian(matrix)
-    assert is_psd_generalized_matrix(gramian)
     reshaped_gramian = reshape(gramian, target_shape)
     assert_psd_generalized_matrix(reshaped_gramian, atol=1e-04, rtol=0.0)
 
@@ -78,7 +76,6 @@ def test_reshape_yields_psd(original_shape: list[int], target_shape: list[int]):
 def test_flatten_yields_matrix(shape: list[int]):
     matrix = randn_(shape + [2])
     gramian = compute_gramian(matrix)
-    assert is_psd_generalized_matrix(gramian)
     flattened_gramian = flatten(gramian)
     assert is_psd_matrix(flattened_gramian)
 
@@ -96,7 +93,6 @@ def test_flatten_yields_matrix(shape: list[int]):
 def test_flatten_yields_psd(shape: list[int]):
     matrix = randn_(shape + [2])
     gramian = compute_gramian(matrix)
-    assert is_psd_generalized_matrix(gramian)
     flattened_gramian = flatten(gramian)
     assert_psd_matrix(flattened_gramian, atol=1e-04, rtol=0.0)
 
@@ -128,7 +124,6 @@ def test_movedim_equivariance(shape: list[int], source: list[int], destination:
     original_gramian = compute_gramian(original_matrix)
     target_gramian = compute_gramian(target_matrix)
 
-    assert is_psd_generalized_matrix(original_gramian)
     moveddim_gramian = movedim(original_gramian, source, destination)
 
     assert_close(moveddim_gramian, target_gramian)
@@ -155,6 +150,5 @@ def test_movedim_equivariance(shape: list[int], source: list[int], destination:
 def test_movedim_yields_psd(shape: list[int], source: list[int], destination: list[int]):
     matrix = randn_(shape + [2])
     gramian = compute_gramian(matrix)
-    assert is_psd_generalized_matrix(gramian)
     moveddim_gramian = movedim(gramian, source, destination)
     assert_psd_generalized_matrix(moveddim_gramian)

tests/utils/forward_backwards.py

@@ -8,6 +8,7 @@
 from utils.architectures import get_in_out_shapes
 from utils.contexts import fork_rng
 
+from torchjd._linalg import PSDGeneralizedMatrix
 from torchjd.aggregation import Aggregator, Weighting
 from torchjd.autogram import Engine
 from torchjd.autojac import backward
@@ -116,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
-) -> Tensor:
+) -> PSDGeneralizedMatrix:
     """
     Computes the Gramian of the Jacobian of the outputs with respect to the params using vmapped
     calls to the autograd engine.
@@ -141,7 +142,7 @@ def get_vjp(grad_outputs: Tensor) -> list[Tensor]:
     return gramian
 
 
-def compute_gramian(matrix: Tensor) -> Tensor:
+def compute_gramian(matrix: Tensor) -> PSDGeneralizedMatrix:
     """Contracts the last dimension of matrix to make it into a Gramian."""
 
     indices = list(range(matrix.ndim))