test: Fix automatic device (#391)

* Remove torch.set_default_device(DEVICE) from conftest.py
* Add curried versions of functions requiring a device in _utils.py
* Use curried fucntions in tests/unit
* Fix device of nn.modules in tests
* Add explanation about how to test on the right device in CONTRIBUTING.md

Author: ValerianRey

CONTRIBUTING.md

@@ -99,6 +99,30 @@ should create it.
 We ask contributors to implement the unit tests necessary to check the correctness of their
 implementations. Besides, whenever usage examples are provided, we require the example's code to be
 tested in `tests/doc`. We require a very high code coverage for newly introduced sources (~95-100%).
+To ensure that the tensors generated during the tests are on the right device, you have to use the
+partial functions defined in `tests/unit/_utils.py` to instantiate tensors. For instance, instead of
+```python
+import torch
+a = torch.ones(3, 4)
+```
+use
+```python
+from unit._utils import ones_
+a = ones_(3, 4)
+```
+
+This will automatically call `torch.ones` with `device=unit.conftest.DEVICE`.
+If the function you need does not exist yet as a partial function in `_utils.py`, add it.
+Lastly, when you create a model or a random generator, you have to move them manually to the right
+device (the `DEVICE` defined in `unit.conftest`):
+```python
+import torch
+from torch.nn import Linear
+from unit.conftest import DEVICE
+
+model = Linear(3, 4).to(device=DEVICE)
+rng = torch.Generator(device=DEVICE)
+```
 
 ### Coding
 

tests/unit/_utils.py

@@ -1,4 +1,22 @@
 from contextlib import AbstractContextManager
+from functools import partial
 from typing import TypeAlias
 
+import torch
+from unit.conftest import DEVICE
+
 ExceptionContext: TypeAlias = AbstractContextManager[Exception | None]
+
+# Curried calls to torch functions that require a device so that we automatically fix the device
+# for code written in the tests, while not affecting code written in src (what
+# torch.set_default_device or what a too large `with torch.device(DEVICE)` context would have done).
+
+empty_ = partial(torch.empty, device=DEVICE)
+eye_ = partial(torch.eye, device=DEVICE)
+ones_ = partial(torch.ones, device=DEVICE)
+rand_ = partial(torch.rand, device=DEVICE)
+randint_ = partial(torch.randint, device=DEVICE)
+randn_ = partial(torch.randn, device=DEVICE)
+randperm_ = partial(torch.randperm, device=DEVICE)
+tensor_ = partial(torch.tensor, device=DEVICE)
+zeros_ = partial(torch.zeros, device=DEVICE)

tests/unit/aggregation/_asserts.py

@@ -2,6 +2,7 @@
 from pytest import raises
 from torch import Tensor
 from torch.testing import assert_close
+from unit._utils import rand_, randperm_
 
 from torchjd.aggregation import Aggregator
 from torchjd.aggregation._utils.non_differentiable import NonDifferentiableError
@@ -44,7 +45,7 @@ def assert_permutation_invariant(
     """
 
     def permute_randomly(matrix_: Tensor) -> Tensor:
-        row_permutation = torch.randperm(matrix_.size(dim=0))
+        row_permutation = randperm_(matrix_.size(dim=0))
         return matrix_[row_permutation]
 
     vector = aggregator(matrix)
@@ -66,10 +67,10 @@ def assert_linear_under_scaling(
     """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)
+        c1 = rand_(matrix.shape[0], dtype=matrix.dtype)
+        c2 = rand_(matrix.shape[0], dtype=matrix.dtype)
+        alpha = rand_([], dtype=matrix.dtype)
+        beta = rand_([], dtype=matrix.dtype)
 
         x1 = aggregator(torch.diag(c1) @ matrix)
         x2 = aggregator(torch.diag(c2) @ matrix)

tests/unit/aggregation/_inputs.py

@@ -1,4 +1,5 @@
 import torch
+from unit._utils import zeros_
 from unit.conftest import DEVICE
 
 from ._matrix_samplers import NonWeakSampler, NormalSampler, StrictlyWeakSampler, StrongSampler
@@ -31,7 +32,7 @@
 _rng = torch.Generator(device=DEVICE).manual_seed(0)
 
 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]
+zero_matrices = [zeros_([m, n]) for m, n, _ in _zero_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]

tests/unit/aggregation/_matrix_samplers.py

@@ -3,6 +3,7 @@
 import torch
 from torch import Tensor
 from torch.nn.functional import normalize
+from unit._utils import randint_, randn_, randperm_, zeros_
 
 
 class MatrixSampler(ABC):
@@ -43,7 +44,7 @@ class NormalSampler(MatrixSampler):
     def __call__(self, rng: torch.Generator | None = None) -> Tensor:
         U = _sample_orthonormal_matrix(self.m, dtype=self.dtype, rng=rng)
         Vt = _sample_orthonormal_matrix(self.n, dtype=self.dtype, rng=rng)
-        S = torch.diag(torch.abs(torch.randn([self.rank], dtype=self.dtype, generator=rng)))
+        S = torch.diag(torch.abs(randn_([self.rank], dtype=self.dtype, generator=rng)))
         A = U[:, : self.rank] @ S @ Vt[: self.rank, :]
         return A
 
@@ -65,11 +66,11 @@ def _check_params(self, m: int, n: int, rank: int, dtype: torch.dtype) -> None:
         assert 0 < rank <= min(m - 1, n)
 
     def __call__(self, rng: torch.Generator | None = None) -> Tensor:
-        v = torch.abs(torch.randn([self.m], dtype=self.dtype, generator=rng))
+        v = torch.abs(randn_([self.m], dtype=self.dtype, generator=rng))
         U1 = normalize(v, dim=0).unsqueeze(1)
         U2 = _sample_semi_orthonormal_complement(U1, rng=rng)
         Vt = _sample_orthonormal_matrix(self.n, dtype=self.dtype, rng=rng)
-        S = torch.diag(torch.abs(torch.randn([self.rank], dtype=self.dtype, generator=rng)))
+        S = torch.diag(torch.abs(randn_([self.rank], dtype=self.dtype, generator=rng)))
         A = U2[:, : self.rank] @ S @ Vt[: self.rank, :]
         return A
 
@@ -101,18 +102,18 @@ def _check_params(self, m: int, n: int, rank: int, dtype: torch.dtype) -> None:
         assert 0 < rank <= min(m - 1, n)
 
     def __call__(self, rng: torch.Generator | None = None) -> Tensor:
-        u = torch.abs(torch.randn([self.m], dtype=self.dtype, generator=rng))
-        split_index = torch.randint(1, self.m, [], generator=rng).item()
-        shuffled_range = torch.randperm(self.m, generator=rng)
-        v = torch.zeros(self.m, dtype=self.dtype)
+        u = torch.abs(randn_([self.m], dtype=self.dtype, generator=rng))
+        split_index = randint_(1, self.m, [], generator=rng).item()
+        shuffled_range = randperm_(self.m, generator=rng)
+        v = zeros_(self.m, dtype=self.dtype)
         v[shuffled_range[:split_index]] = normalize(u[shuffled_range[:split_index]], dim=0)
-        v_prime = torch.zeros(self.m, dtype=self.dtype)
+        v_prime = zeros_(self.m, dtype=self.dtype)
         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, rng=rng)
         U = torch.hstack([U1, U2])
         Vt = _sample_orthonormal_matrix(self.n, dtype=self.dtype, rng=rng)
-        S = torch.diag(torch.abs(torch.randn([self.rank], dtype=self.dtype, generator=rng)))
+        S = torch.diag(torch.abs(randn_([self.rank], dtype=self.dtype, generator=rng)))
         A = U[:, 1 : self.rank + 1] @ S @ Vt[: self.rank, :]
         return A
 
@@ -133,12 +134,12 @@ def _check_params(self, m: int, n: int, rank: int, dtype: torch.dtype) -> None:
         assert 0 < rank
 
     def __call__(self, rng: torch.Generator | None = None) -> Tensor:
-        u = torch.abs(torch.randn([self.m], dtype=self.dtype, generator=rng))
+        u = torch.abs(randn_([self.m], dtype=self.dtype, generator=rng))
         U1 = normalize(u, dim=0).unsqueeze(1)
         U2 = _sample_semi_orthonormal_complement(U1, rng=rng)
         U = torch.hstack([U1, U2])
         Vt = _sample_orthonormal_matrix(self.n, dtype=self.dtype, rng=rng)
-        S = torch.diag(torch.abs(torch.randn([self.rank], dtype=self.dtype, generator=rng)))
+        S = torch.diag(torch.abs(randn_([self.rank], dtype=self.dtype, generator=rng)))
         A = U[:, : self.rank] @ S @ Vt[: self.rank, :]
         return A
 
@@ -148,7 +149,7 @@ def _sample_orthonormal_matrix(
 ) -> Tensor:
     """Uniformly samples a random orthonormal matrix of shape [dim, dim]."""
 
-    return _sample_semi_orthonormal_complement(torch.zeros([dim, 0], dtype=dtype), rng=rng)
+    return _sample_semi_orthonormal_complement(zeros_([dim, 0], dtype=dtype), rng=rng)
 
 
 def _sample_semi_orthonormal_complement(Q: Tensor, rng: torch.Generator | None = None) -> Tensor:
@@ -161,7 +162,7 @@ def _sample_semi_orthonormal_complement(Q: Tensor, rng: torch.Generator | None =
 
     dtype = Q.dtype
     m, k = Q.shape
-    A = torch.randn([m, m - k], dtype=dtype, generator=rng)
+    A = randn_([m, m - k], dtype=dtype, generator=rng)
 
     # project A onto the orthogonal complement of Q
     A_proj = A - Q @ (Q.T @ A)

tests/unit/aggregation/_utils/test_dual_cone.py

@@ -2,6 +2,7 @@
 import torch
 from pytest import mark, raises
 from torch.testing import assert_close
+from unit._utils import rand_, randn_
 
 from torchjd.aggregation._utils.dual_cone import _project_weight_vector, project_weights
 
@@ -29,9 +30,9 @@ def test_solution_weights(shape: tuple[int, int]):
     [1] `Jacobian Descent For Multi-Objective Optimization <https://arxiv.org/pdf/2406.16232>`_.
     """
 
-    J = torch.randn(shape)
+    J = randn_(shape)
     G = J @ J.T
-    u = torch.rand(shape[0])
+    u = rand_(shape[0])
 
     w = project_weights(u, G, "quadprog")
     dual_gap = w - u
@@ -58,9 +59,9 @@ def test_scale_invariant(shape: tuple[int, int], scaling: float):
     Tests that `_project_weights` is invariant under scaling.
     """
 
-    J = torch.randn(shape)
+    J = randn_(shape)
     G = J @ J.T
-    u = torch.rand(shape[0])
+    u = rand_(shape[0])
 
     w = project_weights(u, G, "quadprog")
     w_scaled = project_weights(u, scaling * G, "quadprog")
@@ -75,8 +76,8 @@ def test_tensorization_shape(shape: tuple[int, ...]):
     reshaped as matrix and to reshape the result back to the original tensor's shape.
     """
 
-    matrix = torch.randn([shape[-1], shape[-1]])
-    U_tensor = torch.randn(shape)
+    matrix = randn_([shape[-1], shape[-1]])
+    U_tensor = randn_(shape)
     U_matrix = U_tensor.reshape([-1, shape[-1]])
 
     G = matrix @ matrix.T

tests/unit/aggregation/_utils/test_pref_vector.py

@@ -1,9 +1,8 @@
 from contextlib import nullcontext as does_not_raise
 
-import torch
 from pytest import mark, raises
 from torch import Tensor
-from unit._utils import ExceptionContext
+from unit._utils import ExceptionContext, ones_
 
 from torchjd.aggregation._mean import _MeanWeighting
 from torchjd.aggregation._utils.pref_vector import pref_vector_to_weighting
@@ -13,12 +12,12 @@
     ["pref_vector", "expectation"],
     [
         (None, does_not_raise()),
-        (torch.ones([]), raises(ValueError)),
-        (torch.ones([0]), does_not_raise()),
-        (torch.ones([1]), does_not_raise()),
-        (torch.ones([5]), does_not_raise()),
-        (torch.ones([1, 1]), raises(ValueError)),
-        (torch.ones([1, 1, 1]), raises(ValueError)),
+        (ones_([]), raises(ValueError)),
+        (ones_([0]), does_not_raise()),
+        (ones_([1]), does_not_raise()),
+        (ones_([5]), does_not_raise()),
+        (ones_([1, 1]), raises(ValueError)),
+        (ones_([1, 1, 1]), raises(ValueError)),
     ],
 )
 def test_pref_vector_to_weighting_check(pref_vector: Tensor | None, expectation: ExceptionContext):

tests/unit/aggregation/test_aggregator_bases.py

@@ -1,9 +1,8 @@
 from collections.abc import Sequence
 from contextlib import nullcontext as does_not_raise
 
-import torch
 from pytest import mark, raises
-from unit._utils import ExceptionContext
+from unit._utils import ExceptionContext, randn_
 
 from torchjd.aggregation import Aggregator
 
@@ -20,4 +19,4 @@
 )
 def test_check_is_matrix(shape: Sequence[int], expectation: ExceptionContext):
     with expectation:
-        Aggregator._check_is_matrix(torch.randn(shape))
+        Aggregator._check_is_matrix(randn_(shape))

tests/unit/aggregation/test_cagrad.py

@@ -1,9 +1,8 @@
 from contextlib import nullcontext as does_not_raise
 
-import torch
 from pytest import mark, raises
 from torch import Tensor
-from unit._utils import ExceptionContext
+from unit._utils import ExceptionContext, ones_
 
 from torchjd.aggregation import CAGrad
 
@@ -12,7 +11,7 @@
 
 scaled_pairs = [(CAGrad(c=0.5), matrix) for matrix in scaled_matrices]
 typical_pairs = [(CAGrad(c=0.5), matrix) for matrix in typical_matrices]
-requires_grad_pairs = [(CAGrad(c=0.5), torch.ones(3, 5, requires_grad=True))]
+requires_grad_pairs = [(CAGrad(c=0.5), ones_(3, 5, requires_grad=True))]
 non_conflicting_pairs_1 = [(CAGrad(c=1.0), matrix) for matrix in typical_matrices]
 non_conflicting_pairs_2 = [(CAGrad(c=2.0), matrix) for matrix in typical_matrices]
 

tests/unit/aggregation/test_config.py

@@ -1,6 +1,7 @@
 import torch
 from pytest import mark
 from torch import Tensor
+from unit._utils import ones_
 
 from torchjd.aggregation import ConFIG
 
@@ -15,7 +16,7 @@
 scaled_pairs = [(ConFIG(), matrix) for matrix in scaled_matrices]
 typical_pairs = [(ConFIG(), matrix) for matrix in typical_matrices]
 non_strong_pairs = [(ConFIG(), matrix) for matrix in non_strong_matrices]
-requires_grad_pairs = [(ConFIG(), torch.ones(3, 5, requires_grad=True))]
+requires_grad_pairs = [(ConFIG(), ones_(3, 5, requires_grad=True))]
 
 
 @mark.parametrize(["aggregator", "matrix"], scaled_pairs + typical_pairs)