fix(aggregation): Register tensor attributes as buffers

CHANGELOG.md

@@ -17,6 +17,11 @@ changelog does not include internal changes that do not affect the user.
 - Fixed `Krum` and `KrumWeighting` ranking rows by sums of Euclidean distances instead of
   squared Euclidean distances, as required by Krum and Multi-Krum. This could select a different
   gradient or a different set of gradients to average.
+- Fixed `.to()` not moving the tensors given to `Constant`, `ConstantWeighting`, `UPGrad`,
+  `UPGradWeighting`, `DualProj`, `DualProjWeighting`, `AlignedMTL`, `AlignedMTLWeighting`, `ConFIG`,
+  `SDMGradWeighting` and `CRMOGMWeighting` (`weights`, `pref_vector` and `initial_weights`). These
+  tensors are now registered as non-persistent buffers, so they follow the dtype and device of the
+  aggregator or weighting without being added to its `state_dict`.
 
 ## [0.18.0] - 2026-10-08
 

src/torchjd/aggregation/_aligned_mtl.py

@@ -36,6 +36,7 @@ def __init__(
         scale_mode: SUPPORTED_SCALE_MODE = "min",
     ) -> None:
         super().__init__()
+        self.register_buffer("_pref_vector", None, persistent=False)
         self.pref_vector = pref_vector
         self.scale_mode: SUPPORTED_SCALE_MODE = scale_mode
 

src/torchjd/aggregation/_config.py

@@ -31,6 +31,7 @@ class ConFIG(Aggregator, _NonDifferentiable):
 
     def __init__(self, pref_vector: Tensor | None = None) -> None:
         super().__init__()
+        self.register_buffer("_pref_vector", None, persistent=False)
         self.pref_vector = pref_vector
 
     def forward(self, matrix: Matrix, /) -> Tensor:

src/torchjd/aggregation/_constant.py

@@ -14,6 +14,8 @@ class ConstantWeighting(_MatrixWeighting):
     :param weights: The weights to return at each call.
     """
 
+    weights: Tensor
+
     def __init__(self, weights: Tensor) -> None:
         if weights.dim() != 1:
             raise ValueError(
@@ -22,7 +24,7 @@ def __init__(self, weights: Tensor) -> None:
             )
 
         super().__init__()
-        self.weights = weights
+        self.register_buffer("weights", weights, persistent=False)
 
     def forward(self, matrix: Tensor, /) -> Tensor:
         self._check_matrix_shape(matrix)
@@ -48,11 +50,10 @@ class Constant(WeightedAggregator):
 
     def __init__(self, weights: Tensor) -> None:
         super().__init__(weighting=ConstantWeighting(weights=weights))
-        self._weights = weights
 
     def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(weights={repr(self._weights)})"
+        return f"{self.__class__.__name__}(weights={repr(self.weighting.weights)})"
 
     def __str__(self) -> str:
-        weights_str = vector_to_str(self._weights)
+        weights_str = vector_to_str(self.weighting.weights)
         return f"{self.__class__.__name__}([{weights_str}])"

src/torchjd/aggregation/_cr_mogm.py

@@ -98,13 +98,15 @@ class CRMOGMWeighting(Weighting[_T], Stateful):
             cr_mogm.alpha = 1 - current_lr / initial_lr
     """
 
+    _initial_weights: Tensor | None
+
     def __init__(
         self, weighting: Weighting[_T], alpha: float = 0.9, initial_weights: Tensor | None = None
     ) -> None:
         super().__init__()
         self.weighting = weighting
         self.alpha = alpha
-        self._initial_weights = initial_weights
+        self.register_buffer("_initial_weights", initial_weights, persistent=False)
         self.register_buffer("_lambda", None)
 
     @property

src/torchjd/aggregation/_dualproj.py

@@ -28,6 +28,7 @@ def __init__(
         projector: DualConeProjector | None = None,
     ) -> None:
         super().__init__()
+        self.register_buffer("_pref_vector", None, persistent=False)
         self.pref_vector = pref_vector
         self.projector = projector
 

src/torchjd/aggregation/_sdmgrad.py

@@ -107,6 +107,7 @@ def __init__(
         self.momentum = momentum
         self.n_iter = n_iter
         self.lambda_ = lambda_
+        self.register_buffer("_pref_vector", None, persistent=False)
         self.pref_vector = pref_vector
         self.register_buffer("_w", None)
         self._state_key: tuple[int, torch.dtype, torch.device] | None = None

src/torchjd/aggregation/_upgrad.py

@@ -29,6 +29,7 @@ def __init__(
         projector: DualConeProjector | None = None,
     ) -> None:
         super().__init__()
+        self.register_buffer("_pref_vector", None, persistent=False)
         self.pref_vector = pref_vector
         self.projector = projector
 

tests/unit/aggregation/test_aligned_mtl.py

@@ -66,3 +66,17 @@ def test_float64_small_eigenvalue_is_kept() -> None:
     J = tensor_([[1.0, 0.0], [0.0, 1e-4]], dtype=torch.float64)
     result = AlignedMTL()(J)
     assert_close(result, tensor_([5e-5, 5e-5], dtype=torch.float64))
+
+
+def test_to_moves_pref_vector() -> None:
+    A = AlignedMTL(pref_vector=tensor_([1.0, 2.0], dtype=torch.float32)).to(torch.float64)
+    weighting = A.gramian_weighting.weighting
+    assert isinstance(weighting, ConstantWeighting)
+
+    assert A.pref_vector is not None
+    assert A.pref_vector.dtype == torch.float64
+    assert weighting.weights.dtype == torch.float64
+    buffers = list(A.buffers())
+    assert any(buffer is A.pref_vector for buffer in buffers)
+    assert any(buffer is weighting.weights for buffer in buffers)
+    assert len(A.state_dict()) == 0

tests/unit/aggregation/test_config.py

@@ -1,7 +1,7 @@
 import torch
 from pytest import mark
 from torch import Tensor
-from utils.tensors import ones_
+from utils.tensors import ones_, tensor_
 
 from torchjd.aggregation import ConFIG, ConstantWeighting
 
@@ -56,3 +56,17 @@ def test_pref_vector_setter_updates_value() -> None:
     assert A.pref_vector is new_pref
     assert isinstance(A.weighting, ConstantWeighting)
     assert A.weighting.weights is new_pref
+
+
+def test_to_moves_pref_vector() -> None:
+    A = ConFIG(pref_vector=tensor_([1.0, 2.0], dtype=torch.float32)).to(torch.float64)
+    weighting = A.weighting
+    assert isinstance(weighting, ConstantWeighting)
+
+    assert A.pref_vector is not None
+    assert A.pref_vector.dtype == torch.float64
+    assert weighting.weights.dtype == torch.float64
+    buffers = list(A.buffers())
+    assert any(buffer is A.pref_vector for buffer in buffers)
+    assert any(buffer is weighting.weights for buffer in buffers)
+    assert len(A.state_dict()) == 0

tests/unit/aggregation/test_constant.py

@@ -90,3 +90,11 @@ def test_representations() -> None:
     A = Constant(weights=torch.tensor([1.0, 2.0], device="cpu"))
     assert repr(A) == "Constant(weights=tensor([1., 2.]))"
     assert str(A) == "Constant([1., 2.])"
+
+
+def test_to_moves_weights() -> None:
+    A = Constant(tensor_([1.0, 2.0], dtype=torch.float32)).to(torch.float64)
+
+    assert A.weighting.weights.dtype == torch.float64
+    assert any(buffer is A.weighting.weights for buffer in A.buffers())
+    assert len(A.state_dict()) == 0

tests/unit/aggregation/test_cr_mogm.py

@@ -1,3 +1,4 @@
+import torch
 from pytest import mark, raises
 from torch import Tensor
 from torch.testing import assert_close
@@ -201,3 +202,12 @@ def test_zero_columns() -> None:
     aggregator = WeightedAggregator(CRMOGMWeighting(MeanWeighting()))
     out = aggregator(tensor_([]).reshape(2, 0))
     assert out.shape == (0,)
+
+
+def test_to_moves_initial_weights() -> None:
+    W = CRMOGMWeighting(MeanWeighting(), initial_weights=tensor_([0.5, 0.5], dtype=torch.float32))
+    W = W.to(torch.float64)
+
+    buffers = dict(W.named_buffers())
+    assert buffers["_initial_weights"].dtype == torch.float64
+    assert "_initial_weights" not in W.state_dict()

tests/unit/aggregation/test_dualproj.py

@@ -8,7 +8,7 @@
 import torch
 from pytest import mark
 from torch import Tensor
-from utils.tensors import ones_
+from utils.tensors import ones_, tensor_
 
 from ._asserts import (
     assert_expected_structure,
@@ -89,3 +89,17 @@ def test_projector_setter_updates_value() -> None:
     A.projector = new_projector
     assert A.projector is new_projector
     assert A.gramian_weighting.projector is new_projector
+
+
+def test_to_moves_pref_vector() -> None:
+    A = DualProj(pref_vector=tensor_([1.0, 2.0], dtype=torch.float32)).to(torch.float64)
+    weighting = A.gramian_weighting.weighting
+    assert isinstance(weighting, ConstantWeighting)
+
+    assert A.pref_vector is not None
+    assert A.pref_vector.dtype == torch.float64
+    assert weighting.weights.dtype == torch.float64
+    buffers = list(A.buffers())
+    assert any(buffer is A.pref_vector for buffer in buffers)
+    assert any(buffer is weighting.weights for buffer in buffers)
+    assert len(A.state_dict()) == 0

tests/unit/aggregation/test_sdmgrad.py

@@ -170,3 +170,12 @@ def test_non_differentiable() -> None:
     W = SDMGradWeighting()
     weights = W(A)
     assert not weights.requires_grad
+
+
+def test_to_moves_pref_vector() -> None:
+    W = SDMGradWeighting(pref_vector=tensor_([0.5, 0.5], dtype=torch.float32)).to(torch.float64)
+
+    assert W.pref_vector is not None
+    assert W.pref_vector.dtype == torch.float64
+    assert "_pref_vector" in dict(W.named_buffers())
+    assert "_pref_vector" not in W.state_dict()

tests/unit/aggregation/test_upgrad.py

@@ -8,7 +8,7 @@
 import torch
 from pytest import mark
 from torch import Tensor
-from utils.tensors import ones_
+from utils.tensors import ones_, tensor_
 
 from ._asserts import (
     assert_expected_structure,
@@ -95,3 +95,17 @@ def test_projector_setter_updates_value() -> None:
     A.projector = new_projector
     assert A.projector is new_projector
     assert A.gramian_weighting.projector is new_projector
+
+
+def test_to_moves_pref_vector() -> None:
+    A = UPGrad(pref_vector=tensor_([1.0, 2.0], dtype=torch.float32)).to(torch.float64)
+    weighting = A.gramian_weighting.weighting
+    assert isinstance(weighting, ConstantWeighting)
+
+    assert A.pref_vector is not None
+    assert A.pref_vector.dtype == torch.float64
+    assert weighting.weights.dtype == torch.float64
+    buffers = list(A.buffers())
+    assert any(buffer is A.pref_vector for buffer in buffers)
+    assert any(buffer is weighting.weights for buffer in buffers)
+    assert len(A.state_dict()) == 0