Move asserts to tests/utils and use them in doc tests

Author: ValerianRey

tests/doc/test_backward.py

@@ -3,7 +3,7 @@
 obtained `.jac` field.
 """
 
-from torch.testing import assert_close
+from utils.asserts import assert_jac_close
 
 
 def test_backward():
@@ -18,4 +18,4 @@ def test_backward():
 
     backward([y1, y2])
 
-    assert_close(param.jac, torch.tensor([[-1.0, 1.0], [2.0, 4.0]]), rtol=0.0, atol=1e-04)
+    assert_jac_close(param, torch.tensor([[-1.0, 1.0], [2.0, 4.0]]), rtol=0.0, atol=1e-04)

tests/doc/test_jac_to_grad.py

@@ -3,7 +3,7 @@
 the obtained `.grad` field.
 """
 
-from torch.testing import assert_close
+from utils.asserts import assert_grad_close
 
 
 def test_jac_to_grad():
@@ -19,4 +19,4 @@ def test_jac_to_grad():
     backward([y1, y2])  # param now has a .jac field
     jac_to_grad([param], aggregator=UPGrad())  # param now has a .grad field
 
-    assert_close(param.grad, torch.tensor([0.5000, 2.5000]), rtol=0.0, atol=1e-04)
+    assert_grad_close(param, torch.tensor([0.5000, 2.5000]), rtol=0.0, atol=1e-04)

tests/unit/autojac/_transform/test_accumulate.py

@@ -1,5 +1,5 @@
 from pytest import mark, raises
-from unit.autojac._asserts import assert_grad_close, assert_jac_close
+from utils.asserts import assert_grad_close, assert_jac_close
 from utils.dict_assertions import assert_tensor_dicts_are_close
 from utils.tensors import ones_, tensor_, zeros_
 

tests/unit/autojac/test_backward.py

@@ -1,13 +1,12 @@
 import torch
 from pytest import mark, raises
+from utils.asserts import assert_has_jac, assert_has_no_jac, assert_jac_close
 from utils.tensors import randn_, tensor_
 
 from torchjd.autojac import backward
 from torchjd.autojac._backward import _create_transform
 from torchjd.autojac._transform import OrderedSet
 
-from ._asserts import assert_has_jac, assert_has_no_jac, assert_jac_close
-
 
 def test_check_create_transform():
     """Tests that _create_transform creates a valid Transform."""

tests/unit/autojac/test_jac_to_grad.py

@@ -1,5 +1,5 @@
 from pytest import mark, raises
-from unit.autojac._asserts import assert_grad_close, assert_has_jac, assert_has_no_jac
+from utils.asserts import assert_grad_close, assert_has_jac, assert_has_no_jac
 from utils.tensors import tensor_
 
 from torchjd.aggregation import Aggregator, Mean, PCGrad, UPGrad

tests/unit/autojac/test_mtl_backward.py

@@ -2,20 +2,19 @@
 from pytest import mark, raises
 from settings import DTYPE
 from torch.autograd import grad
-from utils.tensors import arange_, rand_, randn_, tensor_
-
-from torchjd.autojac import mtl_backward
-from torchjd.autojac._mtl_backward import _create_transform
-from torchjd.autojac._transform import OrderedSet
-
-from ._asserts import (
+from utils.asserts import (
     assert_grad_close,
     assert_has_grad,
     assert_has_jac,
     assert_has_no_grad,
     assert_has_no_jac,
     assert_jac_close,
 )
+from utils.tensors import arange_, rand_, randn_, tensor_
+
+from torchjd.autojac import mtl_backward
+from torchjd.autojac._mtl_backward import _create_transform
+from torchjd.autojac._transform import OrderedSet
 
 
 def test_check_create_transform():