Add strong stationarity property to constant, mean, random and sum

Author: PierreQuinton

tests/unit/aggregation/_property_testers.py

@@ -5,11 +5,7 @@
 
 from torchjd.aggregation import Aggregator
 
-from ._inputs import (
-    non_strong_stationary_matrices,
-    scaled_matrices,
-    typical_matrices,
-)
+from ._inputs import non_strong_stationary_matrices, scaled_matrices, typical_matrices
 
 
 class ExpectedStructureProperty:
@@ -122,9 +118,7 @@ def test_stationarity_property(cls, aggregator: Aggregator, matrix: Tensor):
         cls._assert_stationarity_property(aggregator, matrix)
 
     @staticmethod
-    def _assert_stationarity_property(
-        aggregator: Aggregator, matrix: Tensor
-    ) -> None:
+    def _assert_stationarity_property(aggregator: Aggregator, matrix: Tensor) -> None:
         vector = aggregator(matrix)
         norm = vector.norm().item()
         assert norm > 1e-03

tests/unit/aggregation/test_constant.py

@@ -7,8 +7,12 @@
 
 from torchjd.aggregation import Constant
 
-from ._inputs import scaled_matrices, typical_matrices
-from ._property_testers import ExpectedStructureProperty, LinearUnderScalingProperty
+from ._inputs import non_strong_stationary_matrices, scaled_matrices, typical_matrices
+from ._property_testers import (
+    ExpectedStructureProperty,
+    LinearUnderScalingProperty,
+    StrongStationarityProperty,
+)
 
 # The weights must be a vector of length equal to the number of rows in the matrix that it will be
 # applied to. Thus, each `Constant` instance is specific to matrices of a given number of rows. To
@@ -28,8 +32,13 @@ def _make_aggregator(matrix: Tensor) -> Constant:
 _matrices_2 = typical_matrices
 _aggregators_2 = [_make_aggregator(matrix) for matrix in _matrices_2]
 
+_matrices_3 = non_strong_stationary_matrices
+_aggregators_3 = [_make_aggregator(matrix) for matrix in _matrices_3]
+
 
-class TestConstant(ExpectedStructureProperty, LinearUnderScalingProperty):
+class TestConstant(
+    ExpectedStructureProperty, LinearUnderScalingProperty, StrongStationarityProperty
+):
     # Override the parametrization of `test_expected_structure_property` to make the test use the
     # right aggregator with each matrix.
 
@@ -43,6 +52,11 @@ def test_expected_structure_property(cls, aggregator: Constant, matrix: Tensor):
     def test_linear_under_scaling_property(cls, aggregator: Constant, matrix: Tensor):
         cls._assert_linear_under_scaling_property(aggregator, matrix)
 
+    @classmethod
+    @mark.parametrize(["aggregator", "matrix"], zip(_aggregators_3, _matrices_3))
+    def test_stationarity_property(cls, aggregator: Constant, non_stationary_matrix: Tensor):
+        cls._assert_stationarity_property(aggregator, non_stationary_matrix)
+
 
 @mark.parametrize(
     ["weights_shape", "expectation"],

tests/unit/aggregation/test_mean.py

@@ -6,12 +6,16 @@
     ExpectedStructureProperty,
     LinearUnderScalingProperty,
     PermutationInvarianceProperty,
+    StrongStationarityProperty,
 )
 
 
 @mark.parametrize("aggregator", [Mean()])
 class TestMean(
-    ExpectedStructureProperty, PermutationInvarianceProperty, LinearUnderScalingProperty
+    ExpectedStructureProperty,
+    PermutationInvarianceProperty,
+    LinearUnderScalingProperty,
+    StrongStationarityProperty,
 ):
     pass
 

tests/unit/aggregation/test_random.py

@@ -2,11 +2,11 @@
 
 from torchjd.aggregation import Random
 
-from ._property_testers import ExpectedStructureProperty
+from ._property_testers import ExpectedStructureProperty, StrongStationarityProperty
 
 
 @mark.parametrize("aggregator", [Random()])
-class TestRandom(ExpectedStructureProperty):
+class TestRandom(ExpectedStructureProperty, StrongStationarityProperty):
     pass
 
 

tests/unit/aggregation/test_sum.py

@@ -6,11 +6,17 @@
     ExpectedStructureProperty,
     LinearUnderScalingProperty,
     PermutationInvarianceProperty,
+    StrongStationarityProperty,
 )
 
 
 @mark.parametrize("aggregator", [Sum()])
-class TestSum(ExpectedStructureProperty, PermutationInvarianceProperty, LinearUnderScalingProperty):
+class TestSum(
+    ExpectedStructureProperty,
+    PermutationInvarianceProperty,
+    LinearUnderScalingProperty,
+    StrongStationarityProperty,
+):
     pass