Implements the stationarity properties

Author: PierreQuinton

tests/unit/aggregation/_inputs.py

@@ -108,7 +108,7 @@ def _generate_semi_orthonormal_complement(Q: Tensor) -> Tensor:
 strong_stationary_matrices = [
     _generate_strong_stationary_matrix(m, n) for m, n in _stationary_matrices_shapes
 ]
-weak_stationary_matrices = strong_stationary_matrices + [
+weak_stationary_matrices = [
     _generate_weak_stationary_matrix(m, n) for m, n in _stationary_matrices_shapes
 ]
 typical_matrices = zero_matrices + matrices + weak_stationary_matrices + strong_stationary_matrices

tests/unit/aggregation/_property_testers.py

@@ -5,7 +5,13 @@
 
 from torchjd.aggregation import Aggregator
 
-from ._inputs import scaled_matrices, typical_matrices
+from ._inputs import (
+    matrices,
+    scaled_matrices,
+    strong_stationary_matrices,
+    typical_matrices,
+    weak_stationary_matrices,
+)
 
 
 class ExpectedStructureProperty:
@@ -102,3 +108,77 @@ def _assert_linear_under_scaling_property(
         expected = alpha * x1 + beta * x2
 
         assert_close(x, expected, atol=8e-03, rtol=0)
+
+
+class StationarityProperty:
+    """
+    This class tests empirically that a given `Aggregator` satisfies the stationarity property.
+    """
+
+    @staticmethod
+    def _assert_stationarity_property(
+        aggregator: Aggregator,
+        stationary_matrix: Tensor,
+    ) -> None:
+        vector = aggregator(stationary_matrix)
+        norm = vector.norm().item()
+        assert norm < 8e-02
+
+    @staticmethod
+    def _assert_non_stationarity_property(
+        aggregator: Aggregator,
+        non_stationary_matrix: Tensor,
+    ) -> None:
+        vector = aggregator(non_stationary_matrix)
+        norm = vector.norm().item()
+        assert norm > 1e-03
+
+
+class StrongStationarityProperty(StationarityProperty):
+
+    @classmethod
+    @mark.parametrize("stationary_matrix", strong_stationary_matrices)
+    def test_stationarity_property(
+        cls,
+        aggregator: Aggregator,
+        stationary_matrix: Tensor,
+    ):
+        super(StrongStationarityProperty, cls)._assert_stationarity_property(
+            aggregator, stationary_matrix
+        )
+
+    @classmethod
+    @mark.parametrize("non_stationary_matrix", weak_stationary_matrices + matrices)
+    def test_non_stationarity_property(
+        cls,
+        aggregator: Aggregator,
+        non_stationary_matrix: Tensor,
+    ):
+        super(StrongStationarityProperty, cls)._assert_non_stationarity_property(
+            aggregator, non_stationary_matrix
+        )
+
+
+class WeakStationarityProperty(StationarityProperty):
+
+    @classmethod
+    @mark.parametrize("stationary_matrix", strong_stationary_matrices + weak_stationary_matrices)
+    def test_stationarity_property(
+        cls,
+        aggregator: Aggregator,
+        stationary_matrix: Tensor,
+    ):
+        super(WeakStationarityProperty, cls)._assert_stationarity_property(
+            aggregator, stationary_matrix
+        )
+
+    @classmethod
+    @mark.parametrize("non_stationary_matrix", matrices)
+    def test_non_stationarity_property(
+        cls,
+        aggregator: Aggregator,
+        non_stationary_matrix: Tensor,
+    ):
+        super(WeakStationarityProperty, cls)._assert_non_stationarity_property(
+            aggregator, non_stationary_matrix
+        )

tests/unit/aggregation/test_mgda.py

@@ -9,11 +9,17 @@
     ExpectedStructureProperty,
     NonConflictingProperty,
     PermutationInvarianceProperty,
+    WeakStationarityProperty,
 )
 
 
 @mark.parametrize("aggregator", [MGDA()])
-class TestMGDA(ExpectedStructureProperty, NonConflictingProperty, PermutationInvarianceProperty):
+class TestMGDA(
+    ExpectedStructureProperty,
+    NonConflictingProperty,
+    PermutationInvarianceProperty,
+    WeakStationarityProperty,
+):
     pass
 
 

tests/unit/aggregation/test_upgrad.py

@@ -8,6 +8,7 @@
     LinearUnderScalingProperty,
     NonConflictingProperty,
     PermutationInvarianceProperty,
+    StrongStationarityProperty,
 )
 
 
@@ -17,6 +18,7 @@ class TestUPGrad(
     NonConflictingProperty,
     PermutationInvarianceProperty,
     LinearUnderScalingProperty,
+    StrongStationarityProperty,
 ):
     pass