test: Add dtype setting

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/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)

tests/utils/architectures.py

@@ -2,7 +2,7 @@
 
 import torch
 import torchvision
-from device import DEVICE
+from settings import DEVICE, DTYPE
 from torch import Tensor, nn
 from torch.nn import Flatten, ReLU
 from torch.utils._pytree import PyTree
@@ -17,7 +17,7 @@ def __init__(self, architecture: type[nn.Module], *args, **kwargs):
 
     def __call__(self) -> nn.Module:
         with fork_rng(seed=0):
-            return self.architecture(*self.args, **self.kwargs).to(device=DEVICE)
+            return self.architecture(*self.args, **self.kwargs).to(device=DEVICE, dtype=DTYPE)
 
     def __str__(self) -> str:
         args_string = ", ".join([str(arg) for arg in self.args])

tests/utils/contexts.py

@@ -3,7 +3,7 @@
 from typing import Any, TypeAlias
 
 import torch
-from device import DEVICE
+from settings import DEVICE
 
 ExceptionContext: TypeAlias = AbstractContextManager[Exception | None]