improve asserts

Author: PierreQuinton

tests/unit/autogram/test_gramian_utils.py

@@ -1,6 +1,6 @@
 from pytest import mark
 from torch.testing import assert_close
-from utils.asserts import assert_psd_generalized_matrix, assert_psd_matrix
+from utils.asserts import assert_is_psd_generalized_matrix, assert_is_psd_matrix
 from utils.forward_backwards import compute_gramian
 from utils.tensors import randn_
 
@@ -60,7 +60,7 @@ def test_reshape_yields_psd(original_shape: list[int], target_shape: list[int]):
     matrix = randn_(original_shape + [2])
     gramian = compute_gramian(matrix)
     reshaped_gramian = reshape(gramian, target_shape)
-    assert_psd_generalized_matrix(reshaped_gramian, atol=1e-04, rtol=0.0)
+    assert_is_psd_generalized_matrix(reshaped_gramian, atol=1e-04, rtol=0.0)
 
 
 @mark.parametrize(
@@ -94,7 +94,7 @@ def test_flatten_yields_psd(shape: list[int]):
     matrix = randn_(shape + [2])
     gramian = compute_gramian(matrix)
     flattened_gramian = flatten(gramian)
-    assert_psd_matrix(flattened_gramian, atol=1e-04, rtol=0.0)
+    assert_is_psd_matrix(flattened_gramian, atol=1e-04, rtol=0.0)
 
 
 @mark.parametrize(
@@ -151,4 +151,4 @@ def test_movedim_yields_psd(shape: list[int], source: list[int], destination: li
     matrix = randn_(shape + [2])
     gramian = compute_gramian(matrix)
     moveddim_gramian = movedim(gramian, source, destination)
-    assert_psd_generalized_matrix(moveddim_gramian)
+    assert_is_psd_generalized_matrix(moveddim_gramian)

tests/unit/linalg/test_gramian.py

@@ -1,5 +1,5 @@
 from pytest import mark
-from utils.asserts import assert_psd_matrix
+from utils.asserts import assert_is_psd_matrix
 from utils.tensors import randn_
 
 from torchjd._linalg import compute_gramian, is_generalized_matrix, is_matrix
@@ -24,7 +24,7 @@ def test_gramian_is_psd(shape: list[int]):
     matrix = randn_(shape)
     assert is_generalized_matrix(matrix)
     gramian = compute_gramian(matrix)
-    assert_psd_matrix(gramian)
+    assert_is_psd_matrix(gramian)
 
 
 @mark.parametrize(
@@ -42,7 +42,7 @@ def test_normalize_yields_psd(shape: list[int]):
     assert is_matrix(matrix)
     gramian = compute_gramian(matrix)
     normalized_gramian = normalize(gramian, 1e-05)
-    assert_psd_matrix(normalized_gramian)
+    assert_is_psd_matrix(normalized_gramian)
 
 
 @mark.parametrize(
@@ -60,4 +60,4 @@ def test_regularize_yields_psd(shape: list[int]):
     assert is_matrix(matrix)
     gramian = compute_gramian(matrix)
     normalized_gramian = regularize(gramian, 1e-05)
-    assert_psd_matrix(normalized_gramian)
+    assert_is_psd_matrix(normalized_gramian)

tests/utils/asserts.py

@@ -1,39 +1,41 @@
 import torch
+from torch import Tensor
 from torch.testing import assert_close
 
-from torchjd._linalg import PSDGeneralizedMatrix, PSDMatrix
+from torchjd._linalg import is_psd_generalized_matrix, is_psd_matrix
 from torchjd.autogram._gramian_utils import flatten
 from torchjd.autojac._accumulation import is_tensor_with_jac
 
 
-def assert_has_jac(t: torch.Tensor) -> None:
+def assert_has_jac(t: Tensor) -> None:
     assert is_tensor_with_jac(t)
     assert t.jac is not None and t.jac.shape[1:] == t.shape
 
 
-def assert_has_no_jac(t: torch.Tensor) -> None:
+def assert_has_no_jac(t: Tensor) -> None:
     assert not is_tensor_with_jac(t)
 
 
-def assert_jac_close(t: torch.Tensor, expected_jac: torch.Tensor, **kwargs) -> None:
+def assert_jac_close(t: Tensor, expected_jac: Tensor, **kwargs) -> None:
     assert is_tensor_with_jac(t)
     assert_close(t.jac, expected_jac, **kwargs)
 
 
-def assert_has_grad(t: torch.Tensor) -> None:
+def assert_has_grad(t: Tensor) -> None:
     assert (t.grad is not None) and (t.shape == t.grad.shape)
 
 
-def assert_has_no_grad(t: torch.Tensor) -> None:
+def assert_has_no_grad(t: Tensor) -> None:
     assert t.grad is None
 
 
-def assert_grad_close(t: torch.Tensor, expected_grad: torch.Tensor, **kwargs) -> None:
+def assert_grad_close(t: Tensor, expected_grad: Tensor, **kwargs) -> None:
     assert t.grad is not None
     assert_close(t.grad, expected_grad, **kwargs)
 
 
-def assert_psd_matrix(matrix: PSDMatrix, **kwargs) -> None:
+def assert_is_psd_matrix(matrix: Tensor, **kwargs) -> None:
+    assert is_psd_matrix(matrix)
     assert_close(matrix, matrix.mH, **kwargs)
 
     eig_vals = torch.linalg.eigvalsh(matrix)
@@ -42,6 +44,7 @@ def assert_psd_matrix(matrix: PSDMatrix, **kwargs) -> None:
     assert_close(eig_vals, expected_eig_vals, **kwargs)
 
 
-def assert_psd_generalized_matrix(t: PSDGeneralizedMatrix, **kwargs) -> None:
+def assert_is_psd_generalized_matrix(t: Tensor, **kwargs) -> None:
+    assert is_psd_generalized_matrix(t)
     matrix = flatten(t)
-    assert_psd_matrix(matrix, **kwargs)
+    assert_is_psd_matrix(matrix, **kwargs)