diff --git a/CHANGELOG.md b/CHANGELOG.md index dfb2616e..c2d1f1e8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,14 @@ changelog does not include internal changes that do not affect the user. ## [Unreleased] +### Fixed + +- Fixed `ConstantWeighting` returning its weights as-is, regardless of the dtype and device of the + input matrix. `Constant`, as well as `AlignedMTL` and `ConFIG` with a `pref_vector`, failed when + the weights were not of the same dtype or device as the Jacobian, e.g. with a `float32` + `pref_vector` and a `float64` Jacobian. The weights are now cast to the dtype and device of the + input matrix. + ## [0.18.0] - 2026-10-08 ### Added diff --git a/src/torchjd/aggregation/_constant.py b/src/torchjd/aggregation/_constant.py index c849155e..6528661c 100644 --- a/src/torchjd/aggregation/_constant.py +++ b/src/torchjd/aggregation/_constant.py @@ -26,7 +26,7 @@ def __init__(self, weights: Tensor) -> None: def forward(self, matrix: Tensor, /) -> Tensor: self._check_matrix_shape(matrix) - return self.weights + return self.weights.to(dtype=matrix.dtype, device=matrix.device) def _check_matrix_shape(self, matrix: Tensor) -> None: if matrix.shape[0] != len(self.weights): diff --git a/tests/unit/aggregation/test_config.py b/tests/unit/aggregation/test_config.py index fe951640..56b05f0e 100644 --- a/tests/unit/aggregation/test_config.py +++ b/tests/unit/aggregation/test_config.py @@ -1,7 +1,8 @@ import torch from pytest import mark from torch import Tensor -from utils.tensors import ones_ +from torch.testing import assert_close +from utils.tensors import ones_, tensor_ from torchjd.aggregation import ConFIG, ConstantWeighting @@ -49,6 +50,20 @@ def test_representations() -> None: assert str(A) == "ConFIG([1., 2., 3.])" +@mark.parametrize( + ["pref_vector_dtype", "matrix_dtype"], + [(torch.float32, torch.float64), (torch.float64, torch.float32)], +) +def test_pref_vector_dtype_differs_from_matrix_dtype( + pref_vector_dtype: torch.dtype, matrix_dtype: torch.dtype +) -> None: + pref_vector = torch.tensor([1.0, 2.0], device="cpu", dtype=pref_vector_dtype) + matrix = tensor_([[-4.0, 1.0, 1.0], [6.0, 1.0, 1.0]], dtype=matrix_dtype) + expected = ConFIG(pref_vector=pref_vector.to(matrix))(matrix) + + assert_close(ConFIG(pref_vector=pref_vector)(matrix), expected) + + def test_pref_vector_setter_updates_value() -> None: A = ConFIG() new_pref = torch.tensor([1.0, 2.0, 3.0]) diff --git a/tests/unit/aggregation/test_constant.py b/tests/unit/aggregation/test_constant.py index aa1332fc..9b484469 100644 --- a/tests/unit/aggregation/test_constant.py +++ b/tests/unit/aggregation/test_constant.py @@ -3,6 +3,7 @@ import torch from pytest import mark, raises from torch import Tensor +from torch.testing import assert_close from utils.contexts import ExceptionContext from utils.tensors import ones_, tensor_ @@ -86,6 +87,20 @@ def test_matrix_shape_check( _ = aggregator(matrix) +@mark.parametrize( + ["weights_dtype", "matrix_dtype"], + [(torch.float32, torch.float64), (torch.float64, torch.float32)], +) +def test_weights_dtype_differs_from_matrix_dtype( + weights_dtype: torch.dtype, matrix_dtype: torch.dtype +) -> None: + weights = torch.tensor([1.0, 2.0], device="cpu", dtype=weights_dtype) + matrix = tensor_([[-4.0, 1.0, 1.0], [6.0, 1.0, 1.0]], dtype=matrix_dtype) + expected = tensor_([8.0, 3.0, 3.0], dtype=matrix_dtype) + + assert_close(Constant(weights)(matrix), expected) + + def test_representations() -> None: A = Constant(weights=torch.tensor([1.0, 2.0], device="cpu")) assert repr(A) == "Constant(weights=tensor([1., 2.]))"