test: Add dtype setting (#496)

* Rename device.py to settings.py
* Add DTYPE setting through environment variable PYTEST_TORCH_DTYPE
* Make all specifications of float dtype in tests use DTYPE
* Fix test tolerances

Author: ValerianRey

tests/conftest.py

@@ -2,8 +2,8 @@
 from contextlib import nullcontext
 
 import torch
-from device import DEVICE
 from pytest import RaisesExc, fixture, mark
+from settings import DEVICE
 from torch import Tensor
 from utils.architectures import ModuleFactory
 

tests/device.py tests/settings.pytests/device.py renamed to tests/settings.py

@@ -14,3 +14,19 @@
     raise ValueError('Requested device "cuda:0" but cuda is not available.')
 
 DEVICE = torch.device(_device_str)
+
+
+_POSSIBLE_TEST_DTYPES = {"float32", "float64"}
+
+try:
+    _dtype_str = os.environ["PYTEST_TORCH_DTYPE"]
+except KeyError:
+    _dtype_str = "float32"  # Default to float32 if environment variable not set
+
+if _dtype_str not in _POSSIBLE_TEST_DTYPES:
+    raise ValueError(
+        f"Invalid value of environment variable PYTEST_TORCH_DTYPE: {_dtype_str}.\n"
+        f"Possible values: {_POSSIBLE_TEST_DTYPES}."
+    )
+
+DTYPE = getattr(torch, _dtype_str)  # "float32" => torch.float32

tests/speed/autogram/grad_vs_jac_vs_gram.py

@@ -1,7 +1,7 @@
 import gc
 
 import torch
-from device import DEVICE
+from settings import DEVICE
 from utils.architectures import (
     AlexNet,
     Cifar10Model,

tests/unit/aggregation/_inputs.py

@@ -1,5 +1,5 @@
 import torch
-from device import DEVICE
+from settings import DEVICE
 from utils.tensors import zeros_
 
 from ._matrix_samplers import NonWeakSampler, NormalSampler, StrictlyWeakSampler, StrongSampler

tests/unit/aggregation/_matrix_samplers.py

@@ -1,6 +1,7 @@
 from abc import ABC, abstractmethod
 
 import torch
+from settings import DTYPE
 from torch import Tensor
 from torch.nn.functional import normalize
 from utils.tensors import randint_, randn_, randperm_, zeros_
@@ -9,7 +10,7 @@
 class MatrixSampler(ABC):
     """Abstract base class for sampling matrices of a given shape, rank and dtype."""
 
-    def __init__(self, m: int, n: int, rank: int, dtype: torch.dtype = torch.float32):
+    def __init__(self, m: int, n: int, rank: int, dtype: torch.dtype = DTYPE):
         self._check_params(m, n, rank, dtype)
         self.m = m
         self.n = n

tests/unit/aggregation/test_dualproj.py

@@ -27,7 +27,7 @@ def test_expected_structure(aggregator: DualProj, matrix: Tensor):
 
 @mark.parametrize(["aggregator", "matrix"], typical_pairs)
 def test_non_conflicting(aggregator: DualProj, matrix: Tensor):
-    assert_non_conflicting(aggregator, matrix, atol=5e-05, rtol=5e-05)
+    assert_non_conflicting(aggregator, matrix, atol=1e-04, rtol=1e-04)
 
 
 @mark.parametrize(["aggregator", "matrix"], typical_pairs)

tests/unit/aggregation/test_upgrad.py

@@ -28,7 +28,7 @@ def test_expected_structure(aggregator: UPGrad, matrix: Tensor):
 
 @mark.parametrize(["aggregator", "matrix"], typical_pairs)
 def test_non_conflicting(aggregator: UPGrad, matrix: Tensor):
-    assert_non_conflicting(aggregator, matrix, atol=3e-04, rtol=3e-04)
+    assert_non_conflicting(aggregator, matrix, atol=4e-04, rtol=4e-04)
 
 
 @mark.parametrize(["aggregator", "matrix"], typical_pairs)
@@ -38,7 +38,7 @@ def test_permutation_invariant(aggregator: UPGrad, matrix: Tensor):
 
 @mark.parametrize(["aggregator", "matrix"], typical_pairs)
 def test_linear_under_scaling(aggregator: UPGrad, matrix: Tensor):
-    assert_linear_under_scaling(aggregator, matrix, n_runs=5, atol=3e-02, rtol=3e-02)
+    assert_linear_under_scaling(aggregator, matrix, n_runs=5, atol=4e-02, rtol=4e-02)
 
 
 @mark.parametrize(["aggregator", "matrix"], non_strong_pairs)

tests/unit/autojac/_transform/test_aggregate.py

@@ -1,8 +1,8 @@
 import math
 
 import torch
-from device import DEVICE
 from pytest import mark, raises
+from settings import DEVICE
 from utils.dict_assertions import assert_tensor_dicts_are_close
 from utils.tensors import rand_, tensor_, zeros_
 

tests/unit/autojac/test_mtl_backward.py

@@ -1,5 +1,6 @@
 import torch
 from pytest import mark, raises
+from settings import DTYPE
 from torch.autograd import grad
 from torch.testing import assert_close
 from utils.tensors import arange_, rand_, randn_, tensor_
@@ -345,7 +346,7 @@ def test_various_feature_lists(shapes: list[tuple[int]]):
     """Tests that mtl_backward works correctly with various kinds of feature lists."""
 
     p0 = tensor_([1.0, 2.0], requires_grad=True)
-    p1 = arange_(len(shapes), dtype=torch.float32, requires_grad=True)
+    p1 = arange_(len(shapes), dtype=DTYPE, requires_grad=True)
     p2 = tensor_(5.0, requires_grad=True)
 
     features = [rand_(shape) @ p0 for shape in shapes]

tests/unit/autojac/test_utils.py

@@ -1,5 +1,5 @@
-from device import DEVICE
 from pytest import mark, raises
+from settings import DEVICE, DTYPE
 from torch.nn import Linear, MSELoss, ReLU, Sequential
 from utils.tensors import randn_, tensor_
 
@@ -85,7 +85,7 @@ def test_get_leaf_tensors_model():
     x = randn_(16, 10)
     y = randn_(16, 1)
 
-    model = Sequential(Linear(10, 5), ReLU(), Linear(5, 1)).to(device=DEVICE)
+    model = Sequential(Linear(10, 5), ReLU(), Linear(5, 1)).to(device=DEVICE, dtype=DTYPE)
     loss_fn = MSELoss(reduction="none")
 
     y_hat = model(x)
@@ -104,8 +104,8 @@ def test_get_leaf_tensors_model_excluded_2():
     x = randn_(16, 10)
     z = randn_(16, 1)
 
-    model1 = Sequential(Linear(10, 5), ReLU()).to(device=DEVICE)
-    model2 = Linear(5, 1).to(device=DEVICE)
+    model1 = Sequential(Linear(10, 5), ReLU()).to(device=DEVICE, dtype=DTYPE)
+    model2 = Linear(5, 1).to(device=DEVICE, dtype=DTYPE)
     loss_fn = MSELoss(reduction="none")
 
     y = model1(x)