test(aggregation): Refactor property testing (#339)

* Transform property testers in _property_testers.py into assertion functions in _asserts.py
* Add n_runs parameter to assert_linear_under_scaling
* Update aggregator tests to parameterize and use the new assertion functions rather than property testers
* Fix dtype issues
* Add test_permutation_invariant to ConFIG and IMTLG
* Update tolerances

* Refactor matrix sampling functions into classes and move them to _matrix_samplers.py
* Add dtype parameter to the constructor of MatrixSamplers
* Add rng parameter to the call method of MatrixSamplers

Note that we decided to have default values for some parameters of the assertion functions (atol, rtol, n_runs and threshold), so that we still have the benefit of having most aggregators tested with the same tolerances, for ease of maintenance (as we don't want to update each aggregator's tolerance individually each time inputs, hardware or cuda drivers change). We have decided that only UPGrad and DualProj would use custom (and tight) tolerances, so that we're aware if we make a regression for them.

For some reason, the permutation_invariance tests now pass on cuda for ConFIG and IMTLG, so they now test it too. This could be due to a change in the cuda implementation of pinv.

Author: ValerianRey

tests/unit/aggregation/_asserts.py

@@ -0,0 +1,106 @@
+import torch
+from pytest import raises
+from torch import Tensor
+from torch.testing import assert_close
+
+from torchjd.aggregation import Aggregator
+from torchjd.aggregation._utils.non_differentiable import NonDifferentiableError
+
+
+def assert_expected_structure(aggregator: Aggregator, matrix: Tensor) -> None:
+    """
+    Tests that the vector returned by the `__call__` method of an `Aggregator` has the expected
+    structure: it should return a vector whose dimension should be the number of columns of the
+    input matrix, and that should only contain finite values (no `nan`, `inf` or `-inf`). Note that
+    this property implies that the `__call__` method does not raise any exception.
+    """
+
+    vector = aggregator(matrix)  # Will fail if the call raises an exception
+    assert vector.shape == matrix.shape[1:]
+    assert vector.isfinite().all()
+
+
+def assert_non_conflicting(
+    aggregator: Aggregator, matrix: Tensor, atol: float = 4e-04, rtol: float = 4e-04
+) -> None:
+    """Tests empirically that a given `Aggregator` satisfies the non-conflicting property."""
+
+    vector = aggregator(matrix)
+    output_direction = matrix @ vector
+    positive_directions = output_direction[output_direction >= 0]
+    assert_close(positive_directions.norm(), output_direction.norm(), atol=atol, rtol=rtol)
+
+
+def assert_permutation_invariant(
+    aggregator: Aggregator,
+    matrix: Tensor,
+    n_runs: int = 5,
+    atol: float = 1e-04,
+    rtol: float = 1e-04,
+) -> None:
+    """
+    Tests empirically that for a given `Aggregator`, randomly permuting rows of the input matrix
+    doesn't change the aggregation.
+    """
+
+    def permute_randomly(matrix_: Tensor) -> Tensor:
+        row_permutation = torch.randperm(matrix_.size(dim=0))
+        return matrix_[row_permutation]
+
+    vector = aggregator(matrix)
+
+    for _ in range(n_runs):
+        permuted_matrix = permute_randomly(matrix)
+        permuted_vector = aggregator(permuted_matrix)
+
+        assert_close(vector, permuted_vector, atol=atol, rtol=rtol)
+
+
+def assert_linear_under_scaling(
+    aggregator: Aggregator,
+    matrix: Tensor,
+    n_runs: int = 5,
+    atol: float = 1e-04,
+    rtol: float = 1e-04,
+) -> None:
+    """Tests empirically that a given `Aggregator` satisfies the linear under scaling property."""
+
+    for _ in range(n_runs):
+        c1 = torch.rand(matrix.shape[0], dtype=matrix.dtype)
+        c2 = torch.rand(matrix.shape[0], dtype=matrix.dtype)
+        alpha = torch.rand([], dtype=matrix.dtype)
+        beta = torch.rand([], dtype=matrix.dtype)
+
+        x1 = aggregator(torch.diag(c1) @ matrix)
+        x2 = aggregator(torch.diag(c2) @ matrix)
+        x = aggregator(torch.diag(alpha * c1 + beta * c2) @ matrix)
+        expected = alpha * x1 + beta * x2
+
+        assert_close(x, expected, atol=atol, rtol=rtol)
+
+
+def assert_strongly_stationary(
+    aggregator: Aggregator, matrix: Tensor, threshold: float = 5e-03
+) -> None:
+    """
+    Tests empirically that a given `Aggregator` is strongly stationary.
+
+    An aggregator `A` is strongly stationary if for any matrix `J` with `A(J)=0`, `J` is strongly
+    stationary, i.e., there exists `0<w` such that `J^T w=0`. In this class, we test the
+    contraposition: whenever `J` is not strongly stationary, we must have `A(J) != 0`.
+    """
+
+    vector = aggregator(matrix)
+    norm = vector.norm().item()
+    assert norm > threshold
+
+
+def assert_non_differentiable(aggregator: Aggregator, matrix: Tensor):
+    """
+    Tests empirically that a given non-differentiable `Aggregator` correctly raises a
+    NonDifferentiableError whenever we try to backward through it.
+    """
+
+    vector = aggregator(matrix)
+    with raises(NonDifferentiableError):
+        vector.backward(torch.ones_like(vector))

tests/unit/aggregation/_inputs.py

@@ -1,123 +1,7 @@
 import torch
-from torch import Tensor
-from torch.nn.functional import normalize
-
-
-def _sample_matrix(m: int, n: int, rank: int) -> Tensor:
-    """Samples a random matrix A of shape [m, n] with provided rank."""
-
-    U = _sample_orthonormal_matrix(m)
-    Vt = _sample_orthonormal_matrix(n)
-    S = torch.diag(torch.abs(torch.randn([rank])))
-    A = U[:, :rank] @ S @ Vt[:rank, :]
-    return A
-
-
-def _sample_strong_matrix(m: int, n: int, rank: int) -> Tensor:
-    """
-    Samples a random strongly stationary matrix A of shape [m, n] with provided rank.
-
-    Definition: A matrix A is said to be strongly stationary if there exists a vector 0 < v such
-    that v^T A = 0.
-
-    This is done by sampling a positive v, and by then sampling a matrix orthogonal to v.
-    """
-
-    assert 1 < m
-    assert 0 < rank <= min(m - 1, n)
-
-    v = torch.abs(torch.randn([m]))
-    U1 = normalize(v, dim=0).unsqueeze(1)
-    U2 = _sample_semi_orthonormal_complement(U1)
-    Vt = _sample_orthonormal_matrix(n)
-    S = torch.diag(torch.abs(torch.randn([rank])))
-    A = U2[:, :rank] @ S @ Vt[:rank, :]
-    return A
-
-
-def _sample_strictly_weak_matrix(m: int, n: int, rank: int) -> Tensor:
-    """
-    Samples a random strictly weakly stationary matrix A of shape [m, n] with provided rank.
-
-    Definition: A matrix A is said to be weakly stationary if there exists a vector 0 <= v, v != 0,
-    such that v^T A = 0.
-
-    Definition: A matrix A is said to be strictly weakly stationary if it is weakly stationary and
-    not strongly stationary, i.e. if there exists a vector 0 <= v, v != 0, such that v^T A = 0 and
-    there exists no vector 0 < w with w^T A = 0.
-
-    This is done by sampling two unit-norm vectors v, v', whose sum u is a positive vector. These
-    two vectors are also non-negative and non-zero, and are furthermore orthogonal. Then, a matrix
-    A, orthogonal to v, is sampled. By its orthogonality to v, A is weakly stationary. Moreover,
-    since v' is a non-negative left-singular vector of A with positive singular value s, any 0 < w
-    satisfies w^T A != 0. Otherwise, we would have 0 = w^T A A^T v' = s w^T v' > 0, which is a
-    contradiction. A is thus also not strongly stationary.
-    """
-
-    assert 1 < m
-    assert 0 < rank <= min(m - 1, n)
-
-    u = torch.abs(torch.randn([m]))
-    split_index = torch.randint(1, m, []).item()
-    shuffled_range = torch.randperm(m)
-    v = torch.zeros(m)
-    v[shuffled_range[:split_index]] = normalize(u[shuffled_range[:split_index]], dim=0)
-    v_prime = torch.zeros(m)
-    v_prime[shuffled_range[split_index:]] = normalize(u[shuffled_range[split_index:]], dim=0)
-    U1 = torch.stack([v, v_prime]).T
-    U2 = _sample_semi_orthonormal_complement(U1)
-    U = torch.hstack([U1, U2])
-    Vt = _sample_orthonormal_matrix(n)
-    S = torch.diag(torch.abs(torch.randn([rank])))
-    A = U[:, 1 : rank + 1] @ S @ Vt[:rank, :]
-    return A
-
-
-def _sample_non_weak_matrix(m: int, n: int, rank: int) -> Tensor:
-    """
-    Samples a random non weakly-stationary matrix A of shape [m, n] with provided rank.
-
-    This is done by sampling a positive u, and by then sampling a matrix A that has u as one of its
-    left-singular vectors, with positive singular value s. Any 0 <= v, v != 0, satisfies v^T A != 0.
-    Otherwise, we would have 0 = v^T A A^T u = s v^T u > 0, which is a contradiction. A is thus not
-    weakly stationary.
-    """
-
-    assert 0 < rank <= min(m, n)
-
-    u = torch.abs(torch.randn([m]))
-    U1 = normalize(u, dim=0).unsqueeze(1)
-    U2 = _sample_semi_orthonormal_complement(U1)
-    U = torch.hstack([U1, U2])
-    Vt = _sample_orthonormal_matrix(n)
-    S = torch.diag(torch.abs(torch.randn([rank])))
-    A = U[:, :rank] @ S @ Vt[:rank, :]
-    return A
-
-
-def _sample_orthonormal_matrix(dim: int) -> Tensor:
-    """Uniformly samples a random orthonormal matrix of shape [dim, dim]."""
-
-    return _sample_semi_orthonormal_complement(torch.zeros([dim, 0]))
-
-
-def _sample_semi_orthonormal_complement(Q: Tensor) -> Tensor:
-    """
-    Uniformly samples a random semi-orthonormal matrix Q' (i.e. Q'^T Q' = I) of shape [m, m-k]
-    orthogonal to Q, i.e. such that the concatenation [Q, Q'] is an orthonormal matrix.
-
-    :param Q: A semi-orthonormal matrix (i.e. Q^T Q = I) of shape [m, k], with k <= m.
-    """
-
-    m, k = Q.shape
-    A = torch.randn([m, m - k])
-
-    # project A onto the orthogonal complement of Q
-    A_proj = A - Q @ (Q.T @ A)
-
-    Q_prime, _ = torch.linalg.qr(A_proj)
-    return Q_prime
+from unit.conftest import DEVICE
 
+from ._matrix_samplers import NonWeakSampler, NormalSampler, StrictlyWeakSampler, StrongSampler
 
 _normal_dims = [
     (1, 1, 1),
@@ -144,14 +28,13 @@ def _sample_semi_orthonormal_complement(Q: Tensor) -> Tensor:
 
 _scales = [0.0, 1e-10, 1e3, 1e5, 1e10, 1e15]
 
-# Fix seed to fix randomness of matrix sampling
-torch.manual_seed(0)
+_rng = torch.Generator(device=DEVICE).manual_seed(0)
 
-matrices = [_sample_matrix(m, n, r) for m, n, r in _normal_dims]
+matrices = [NormalSampler(m, n, r)(_rng) for m, n, r in _normal_dims]
 zero_matrices = [torch.zeros([m, n]) for m, n, _ in _zero_dims]
-strong_matrices = [_sample_strong_matrix(m, n, r) for m, n, r in _stationarity_dims]
-strictly_weak_matrices = [_sample_strictly_weak_matrix(m, n, r) for m, n, r in _stationarity_dims]
-non_weak_matrices = [_sample_non_weak_matrix(m, n, r) for m, n, r in _stationarity_dims]
+strong_matrices = [StrongSampler(m, n, r)(_rng) for m, n, r in _stationarity_dims]
+strictly_weak_matrices = [StrictlyWeakSampler(m, n, r)(_rng) for m, n, r in _stationarity_dims]
+non_weak_matrices = [NonWeakSampler(m, n, r)(_rng) for m, n, r in _stationarity_dims]
 
 scaled_matrices = [scale * matrix for scale in _scales for matrix in matrices]
 
@@ -170,4 +53,4 @@ def _sample_semi_orthonormal_complement(Q: Tensor) -> Tensor:
     (9, 11, 5),
     (9, 11, 9),
 ]
-nash_mtl_matrices = [_sample_matrix(m, n, r) for m, n, r in _nashmtl_dims]
+nash_mtl_matrices = [NormalSampler(m, n, r)(_rng) for m, n, r in _nashmtl_dims]