fix(aggregation): Cast ConstantWeighting weights to matrix dtype

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

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

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

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.]))"