Fix parametrization for aggregators overriding matrix parametrization

Author: ValerianRey

tests/unit/aggregation/_inputs.py

@@ -113,6 +113,5 @@ def _generate_semi_orthonormal_complement(Q: Tensor) -> Tensor:
 ]
 typical_matrices = zero_matrices + matrices + weak_stationary_matrices + strong_stationary_matrices
 
-matrices_2_plus_rows = [matrix for matrix in matrices + zero_matrices if matrix.shape[0] >= 2]
 scaled_matrices_2_plus_rows = [matrix for matrix in scaled_matrices if matrix.shape[0] >= 2]
 typical_matrices_2_plus_rows = [matrix for matrix in typical_matrices if matrix.shape[0] >= 2]

tests/unit/aggregation/test_cagrad.py

@@ -4,7 +4,7 @@
 
 from torchjd.aggregation import CAGrad, Mean
 
-from ._inputs import matrices, strong_stationary_matrices
+from ._inputs import typical_matrices
 from ._property_testers import ExpectedStructureProperty, NonConflictingProperty
 
 
@@ -20,7 +20,7 @@ class TestCAGradNonConflicting(NonConflictingProperty):
     pass
 
 
-@mark.parametrize("matrix", strong_stationary_matrices + matrices)
+@mark.parametrize("matrix", typical_matrices)
 def test_equivalence_mean(matrix: Tensor):
     """Tests that CAGrad is equivalent to Mean when c=0."""
 

tests/unit/aggregation/test_constant.py

@@ -4,7 +4,7 @@
 
 from torchjd.aggregation import Constant
 
-from ._inputs import matrices, scaled_matrices, strong_stationary_matrices, zero_matrices
+from ._inputs import scaled_matrices, typical_matrices
 from ._property_testers import ExpectedStructureProperty, LinearUnderScalingProperty
 
 # The weights must be a vector of length equal to the number of rows in the matrix that it will be
@@ -19,10 +19,10 @@ def _make_aggregator(matrix: Tensor) -> Constant:
     return Constant(weights)
 
 
-_matrices_1 = scaled_matrices + zero_matrices
+_matrices_1 = scaled_matrices + typical_matrices
 _aggregators_1 = [_make_aggregator(matrix) for matrix in _matrices_1]
 
-_matrices_2 = matrices + strong_stationary_matrices
+_matrices_2 = typical_matrices
 _aggregators_2 = [_make_aggregator(matrix) for matrix in _matrices_2]
 
 

tests/unit/aggregation/test_krum.py

@@ -3,7 +3,7 @@
 
 from torchjd.aggregation import Krum
 
-from ._inputs import scaled_matrices_2_plus_rows
+from ._inputs import scaled_matrices_2_plus_rows, typical_matrices_2_plus_rows
 from ._property_testers import ExpectedStructureProperty
 
 
@@ -12,7 +12,7 @@ class TestKrum(ExpectedStructureProperty):
     # Override the parametrization of some property-testing methods because Krum only works on
     # matrices with >= 2 rows.
     @classmethod
-    @mark.parametrize("matrix", scaled_matrices_2_plus_rows)
+    @mark.parametrize("matrix", scaled_matrices_2_plus_rows + typical_matrices_2_plus_rows)
     def test_expected_structure_property(cls, aggregator: Krum, matrix: Tensor):
         cls._assert_expected_structure_property(aggregator, matrix)
 

tests/unit/aggregation/test_trimmed_mean.py

@@ -3,7 +3,7 @@
 
 from torchjd.aggregation import Aggregator, TrimmedMean
 
-from ._inputs import matrices_2_plus_rows, scaled_matrices_2_plus_rows
+from ._inputs import scaled_matrices_2_plus_rows, typical_matrices_2_plus_rows
 from ._property_testers import ExpectedStructureProperty, PermutationInvarianceProperty
 
 
@@ -12,12 +12,12 @@ class TestTrimmedMean(ExpectedStructureProperty, PermutationInvarianceProperty):
     # Override the parametrization of some property-testing methods because `TrimmedMean` with
     # `trim_number=1` only works on matrices with >= 2 rows.
     @classmethod
-    @mark.parametrize("matrix", scaled_matrices_2_plus_rows)
+    @mark.parametrize("matrix", scaled_matrices_2_plus_rows + typical_matrices_2_plus_rows)
     def test_expected_structure_property(cls, aggregator: TrimmedMean, matrix: Tensor):
         cls._assert_expected_structure_property(aggregator, matrix)
 
     @classmethod
-    @mark.parametrize("matrix", matrices_2_plus_rows)
+    @mark.parametrize("matrix", typical_matrices_2_plus_rows)
     def test_permutation_invariance_property(cls, aggregator: Aggregator, matrix: Tensor):
         cls._assert_permutation_invariance_property(aggregator, matrix)