Fix tests

- Remove test_aggregate.py
- Update test_accumulate.py and test_interactions.py to test on AccumulateGrad instead of Accumulate
- Fix tests in test_backward.py and test_mtl_backward.py to match the new interface: check the jac field instead of the .grad field.
- Use _asserts.py for helper functions common to backward.py and mtl_backward.py

Author: ValerianRey

tests/unit/autojac/_asserts.py

@@ -0,0 +1,35 @@
+from typing import cast
+
+import torch
+from torch.testing import assert_close
+
+from torchjd.utils._tensor_with_jac import TensorWithJac
+
+
+def assert_has_jac(t: torch.Tensor) -> None:
+    assert hasattr(t, "jac")
+    t_ = cast(TensorWithJac, t)
+    assert t_.jac is not None and t_.jac.shape[1:] == t_.shape
+
+
+def assert_has_no_jac(t: torch.Tensor) -> None:
+    assert not hasattr(t, "jac")
+
+
+def assert_jac_close(t: torch.Tensor, expected_jac: torch.Tensor) -> None:
+    assert hasattr(t, "jac")
+    t_ = cast(TensorWithJac, t)
+    assert_close(t_.jac, expected_jac)
+
+
+def assert_has_grad(t: torch.Tensor) -> None:
+    assert (t.grad is not None) and (t.shape == t.grad.shape)
+
+
+def assert_has_no_grad(t: torch.Tensor) -> None:
+    assert t.grad is None
+
+
+def assert_grad_close(t: torch.Tensor, expected_grad: torch.Tensor) -> None:
+    assert t.grad is not None
+    assert_close(t.grad, expected_grad)

tests/unit/autojac/_transform/test_accumulate.py

@@ -2,12 +2,12 @@
 from utils.dict_assertions import assert_tensor_dicts_are_close
 from utils.tensors import ones_, tensor_, zeros_
 
-from torchjd.autojac._transform import Accumulate
+from torchjd.autojac._transform import AccumulateGrad
 
 
 def test_single_accumulation():
     """
-    Tests that the Accumulate transform correctly accumulates gradients in .grad fields when run
+    Tests that the AccumulateGrad transform correctly accumulates gradients in .grad fields when run
     once.
     """
 
@@ -19,7 +19,7 @@ def test_single_accumulation():
     value3 = ones_([2, 3])
     input = {key1: value1, key2: value2, key3: value3}
 
-    accumulate = Accumulate()
+    accumulate = AccumulateGrad()
 
     output = accumulate(input)
     expected_output = {}
@@ -35,7 +35,7 @@ def test_single_accumulation():
 @mark.parametrize("iterations", [1, 2, 4, 10, 13])
 def test_multiple_accumulation(iterations: int):
     """
-    Tests that the Accumulate transform correctly accumulates gradients in .grad fields when run
+    Tests that the AccumulateGrad transform correctly accumulates gradients in .grad fields when run
     `iterations` times.
     """
 
@@ -46,7 +46,7 @@ def test_multiple_accumulation(iterations: int):
     value2 = ones_([1])
     value3 = ones_([2, 3])
 
-    accumulate = Accumulate()
+    accumulate = AccumulateGrad()
 
     for i in range(iterations):
         # Clone values to ensure that we accumulate values that are not ever used afterwards
@@ -65,31 +65,31 @@ def test_multiple_accumulation(iterations: int):
 
 def test_no_requires_grad_fails():
     """
-    Tests that the Accumulate transform raises an error when it tries to populate a .grad of a
+    Tests that the AccumulateGrad transform raises an error when it tries to populate a .grad of a
     tensor that does not require grad.
     """
 
     key = zeros_([1], requires_grad=False)
     value = ones_([1])
     input = {key: value}
 
-    accumulate = Accumulate()
+    accumulate = AccumulateGrad()
 
     with raises(ValueError):
         accumulate(input)
 
 
 def test_no_leaf_and_no_retains_grad_fails():
     """
-    Tests that the Accumulate transform raises an error when it tries to populate a .grad of a
+    Tests that the AccumulateGrad transform raises an error when it tries to populate a .grad of a
     tensor that is not a leaf and that does not retain grad.
     """
 
     key = tensor_([1.0], requires_grad=True) * 2
     value = ones_([1])
     input = {key: value}
 
-    accumulate = Accumulate()
+    accumulate = AccumulateGrad()
 
     with raises(ValueError):
         accumulate(input)
@@ -99,7 +99,7 @@ def test_check_keys():
     """Tests that the `check_keys` method works correctly."""
 
     key = tensor_([1.0], requires_grad=True)
-    accumulate = Accumulate()
+    accumulate = AccumulateGrad()
 
     output_keys = accumulate.check_keys({key})
     assert output_keys == set()

tests/unit/autojac/_transform/test_aggregate.py

@@ -5,7 +5,7 @@
 from utils.tensors import tensor_, zeros_
 
 from torchjd.autojac._transform import (
-    Accumulate,
+    AccumulateGrad,
     Conjunction,
     Diagonalize,
     Grad,
@@ -186,18 +186,18 @@ def test_conjunction_is_associative():
 
 def test_conjunction_accumulate_select():
     """
-    Tests that it is possible to conjunct an Accumulate and a Select in this order.
-    It is not trivial since the type of the TensorDict returned by the first transform (Accumulate)
-    is EmptyDict, which is not the type that the conjunction should return (Gradients), but a
-    subclass of it.
+    Tests that it is possible to conjunct an AccumulateGrad and a Select in this order.
+    It is not trivial since the type of the TensorDict returned by the first transform
+    (AccumulateGrad) is EmptyDict, which is not the type that the conjunction should return
+    (Gradients), but a subclass of it.
     """
 
     key = tensor_([1.0, 2.0, 3.0], requires_grad=True)
     value = torch.ones_like(key)
     input = {key: value}
 
     select = Select(set())
-    accumulate = Accumulate()
+    accumulate = AccumulateGrad()
     conjunction = accumulate | select
 
     output = conjunction(input)