Add test of stateful/stateless (all must be one or the other).

Author: PierreQuinton

tests/unit/aggregation/_asserts.py

@@ -1,10 +1,11 @@
 import torch
+from numpy.ma.testutils import assert_allclose
 from pytest import raises
 from torch import Tensor
 from torch.testing import assert_close
 from utils.tensors import rand_, randperm_
 
-from torchjd.aggregation import Aggregator
+from torchjd.aggregation import Aggregator, Stateful
 from torchjd.aggregation._utils.non_differentiable import NonDifferentiableError
 
 
@@ -110,3 +111,33 @@ def assert_non_differentiable(aggregator: Aggregator, matrix: Tensor) -> None:
     vector = aggregator(matrix)
     with raises(NonDifferentiableError):
         vector.backward(torch.ones_like(vector))
+
+
+def assert_stateful(aggregator: Aggregator, matrix: Tensor) -> None:
+    """
+    Test that a given `Aggregator` is stateful. Specifically:
+    - For a fixed state, the aggregator is determinist on the matrix
+    - The reset method and the constructor both set the state to the initial state
+    """
+
+    assert isinstance(aggregator, Stateful)
+
+    first_pair = (aggregator(matrix), aggregator(matrix))
+    aggregator.reset()
+    second_pair = (aggregator(matrix), aggregator(matrix))
+
+    assert_allclose(first_pair[0], second_pair[0], atol=0.0, rtol=0.0)
+    assert_allclose(first_pair[1], second_pair[1], atol=0.0, rtol=0.0)
+
+
+def assert_stateless(aggregator: Aggregator, matrix: Tensor) -> None:
+    """
+    Test that a given `Aggregator` is stateless. Specifically, it must be deterministic.
+    """
+
+    assert not isinstance(aggregator, Stateful)
+
+    first = aggregator(matrix)
+    second = aggregator(matrix)
+
+    assert_allclose(first, second, atol=0.0, rtol=0.0)

tests/unit/aggregation/test_aligned_mtl.py

@@ -5,7 +5,7 @@
 
 from torchjd.aggregation import AlignedMTL
 
-from ._asserts import assert_expected_structure, assert_permutation_invariant
+from ._asserts import assert_expected_structure, assert_permutation_invariant, assert_stateless
 from ._inputs import scaled_matrices, typical_matrices
 
 aggregators = [
@@ -28,6 +28,11 @@ def test_permutation_invariant(aggregator: AlignedMTL, matrix: Tensor) -> None:
     assert_permutation_invariant(aggregator, matrix)
 
 
+@mark.parametrize(["aggregator", "matrix"], typical_pairs)
+def test_stateless(aggregator: AlignedMTL, matrix: Tensor) -> None:
+    assert_stateless(aggregator, matrix)
+
+
 def test_representations() -> None:
     A = AlignedMTL(pref_vector=None)
     assert repr(A) == "AlignedMTL(pref_vector=None, scale_mode='min')"

tests/unit/aggregation/test_cagrad.py

@@ -12,7 +12,12 @@
 
     pytest.skip("CAGrad dependencies not installed", allow_module_level=True)
 
-from ._asserts import assert_expected_structure, assert_non_conflicting, assert_non_differentiable
+from ._asserts import (
+    assert_expected_structure,
+    assert_non_conflicting,
+    assert_non_differentiable,
+    assert_stateless,
+)
 from ._inputs import scaled_matrices, typical_matrices
 
 scaled_pairs = [(CAGrad(c=0.5), matrix) for matrix in scaled_matrices]
@@ -38,6 +43,11 @@ def test_non_conflicting(aggregator: CAGrad, matrix: Tensor) -> None:
     assert_non_conflicting(aggregator, matrix)
 
 
+@mark.parametrize(["aggregator", "matrix"], typical_pairs)
+def test_stateless(aggregator: CAGrad, matrix: Tensor) -> None:
+    assert_stateless(aggregator, matrix)
+
+
 @mark.parametrize(
     ["c", "expectation"],
     [

tests/unit/aggregation/test_config.py

@@ -10,6 +10,7 @@
     assert_linear_under_scaling,
     assert_non_differentiable,
     assert_permutation_invariant,
+    assert_stateless,
 )
 from ._inputs import non_strong_matrices, scaled_matrices, typical_matrices
 
@@ -39,6 +40,11 @@ def test_non_differentiable(aggregator: ConFIG, matrix: Tensor) -> None:
     assert_non_differentiable(aggregator, matrix)
 
 
+@mark.parametrize(["aggregator", "matrix"], typical_pairs)
+def test_stateless(aggregator: ConFIG, matrix: Tensor) -> None:
+    assert_stateless(aggregator, matrix)
+
+
 def test_representations() -> None:
     A = ConFIG()
     assert repr(A) == "ConFIG(pref_vector=None)"

tests/unit/aggregation/test_constant.py

@@ -11,6 +11,7 @@
 from ._asserts import (
     assert_expected_structure,
     assert_linear_under_scaling,
+    assert_stateless,
     assert_strongly_stationary,
 )
 from ._inputs import non_strong_matrices, scaled_matrices, typical_matrices
@@ -42,6 +43,11 @@ def test_strongly_stationary(aggregator: Constant, matrix: Tensor) -> None:
     assert_strongly_stationary(aggregator, matrix)
 
 
+@mark.parametrize(["aggregator", "matrix"], typical_pairs)
+def test_stateless(aggregator: Constant, matrix: Tensor) -> None:
+    assert_stateless(aggregator, matrix)
+
+
 @mark.parametrize(
     ["weights_shape", "expectation"],
     [

tests/unit/aggregation/test_dualproj.py

@@ -10,6 +10,7 @@
     assert_non_conflicting,
     assert_non_differentiable,
     assert_permutation_invariant,
+    assert_stateless,
     assert_strongly_stationary,
 )
 from ._inputs import non_strong_matrices, scaled_matrices, typical_matrices
@@ -45,6 +46,11 @@ def test_non_differentiable(aggregator: DualProj, matrix: Tensor) -> None:
     assert_non_differentiable(aggregator, matrix)
 
 
+@mark.parametrize(["aggregator", "matrix"], typical_pairs)
+def test_stateless(aggregator: DualProj, matrix: Tensor) -> None:
+    assert_stateless(aggregator, matrix)
+
+
 def test_representations() -> None:
     A = DualProj(pref_vector=None, norm_eps=0.0001, reg_eps=0.0001, solver="quadprog")
     assert (

tests/unit/aggregation/test_graddrop.py

@@ -9,7 +9,7 @@
 
 from torchjd.aggregation import GradDrop
 
-from ._asserts import assert_expected_structure, assert_non_differentiable
+from ._asserts import assert_expected_structure, assert_non_differentiable, assert_stateful
 from ._inputs import scaled_matrices, typical_matrices
 
 scaled_pairs = [(GradDrop(), matrix) for matrix in scaled_matrices]
@@ -27,6 +27,11 @@ def test_non_differentiable(aggregator: GradDrop, matrix: Tensor) -> None:
     assert_non_differentiable(aggregator, matrix)
 
 
+@mark.parametrize(["aggregator", "matrix"], typical_pairs)
+def test_stateful(aggregator: GradDrop, matrix: Tensor) -> None:
+    assert_stateful(aggregator, matrix)
+
+
 @mark.parametrize(
     ["leak_shape", "expectation"],
     [

tests/unit/aggregation/test_gradvac.py

@@ -6,7 +6,7 @@
 
 from torchjd.aggregation import GradVac, GradVacWeighting
 
-from ._asserts import assert_expected_structure, assert_non_differentiable
+from ._asserts import assert_expected_structure, assert_non_differentiable, assert_stateful
 from ._inputs import scaled_matrices, typical_matrices, typical_matrices_2_plus_rows
 
 scaled_pairs = [(GradVac(), m) for m in scaled_matrices]
@@ -104,6 +104,11 @@ def test_non_differentiable(aggregator: GradVac, matrix: Tensor) -> None:
     assert_non_differentiable(aggregator, matrix)
 
 
+@mark.parametrize(["aggregator", "matrix"], typical_pairs)
+def test_stateful(aggregator: GradVac, matrix: Tensor) -> None:
+    assert_stateful(aggregator, matrix)
+
+
 def test_weighting_beta_out_of_range() -> None:
     with raises(ValueError, match="beta"):
         GradVacWeighting(beta=-0.1)

tests/unit/aggregation/test_imtl_g.py

@@ -9,6 +9,7 @@
     assert_expected_structure,
     assert_non_differentiable,
     assert_permutation_invariant,
+    assert_stateless,
 )
 from ._inputs import scaled_matrices, typical_matrices
 
@@ -32,6 +33,11 @@ def test_non_differentiable(aggregator: IMTLG, matrix: Tensor) -> None:
     assert_non_differentiable(aggregator, matrix)
 
 
+@mark.parametrize(["aggregator", "matrix"], typical_pairs)
+def test_stateless(aggregator: IMTLG, matrix: Tensor) -> None:
+    assert_stateless(aggregator, matrix)
+
+
 def test_imtlg_zero() -> None:
     """
     Tests that IMTLG correctly returns the 0 vector in the special case where input matrix only

tests/unit/aggregation/test_krum.py

@@ -7,7 +7,7 @@
 
 from torchjd.aggregation import Krum
 
-from ._asserts import assert_expected_structure
+from ._asserts import assert_expected_structure, assert_stateless
 from ._inputs import scaled_matrices_2_plus_rows, typical_matrices_2_plus_rows
 
 scaled_pairs = [(Krum(n_byzantine=1), matrix) for matrix in scaled_matrices_2_plus_rows]
@@ -19,6 +19,11 @@ def test_expected_structure(aggregator: Krum, matrix: Tensor) -> None:
     assert_expected_structure(aggregator, matrix)
 
 
+@mark.parametrize(["aggregator", "matrix"], typical_pairs)
+def test_stateless(aggregator: Krum, matrix: Tensor) -> None:
+    assert_stateless(aggregator, matrix)
+
+
 @mark.parametrize(
     ["n_byzantine", "expectation"],
     [