fix(scalarization): Register tensor attributes as buffers

CHANGELOG.md

@@ -10,6 +10,10 @@ changelog does not include internal changes that do not affect the user.
 
 ### Fixed
 
+- Fixed `.to()` not moving the tensors given to the `Constant`, `STCH`, `PBI` and `COSMOS`
+  scalarizers (`weights` and `reference`). These tensors are now registered as non-persistent
+  buffers, so they follow the dtype and device of the scalarizer without being added to its
+  `state_dict`. `FAMO`'s `min_losses` buffer is also no longer in its `state_dict`.
 - Fixed `.to()` not moving the tensors given to `Constant`, `ConstantWeighting`, `UPGrad`,
   `UPGradWeighting`, `DualProj`, `DualProjWeighting`, `AlignedMTL`, `AlignedMTLWeighting`, `ConFIG`,
   `SDMGradWeighting`, `CRMOGMWeighting` and `GradDrop` (`weights`, `pref_vector`,

src/torchjd/scalarization/_constant.py

@@ -14,9 +14,11 @@ class Constant(Scalarizer):
         passed at call time.
     """
 
+    weights: Tensor
+
     def __init__(self, weights: Tensor) -> None:
         super().__init__()
-        self.weights = weights
+        self.register_buffer("weights", weights, persistent=False)
 
     def forward(self, values: Tensor, /) -> Tensor:
         if values.shape != self.weights.shape:

src/torchjd/scalarization/_cosmos.py

@@ -42,6 +42,8 @@ class COSMOS(Scalarizer):
         objective.
     """
 
+    weights: Tensor
+
     def __init__(self, lambda_: float, weights: Tensor) -> None:
         if lambda_ < 0.0:
             raise ValueError(
@@ -50,7 +52,7 @@ def __init__(self, lambda_: float, weights: Tensor) -> None:
 
         super().__init__()
         self.lambda_ = lambda_
-        self.weights = weights
+        self.register_buffer("weights", weights, persistent=False)
 
     def forward(self, values: Tensor, /) -> Tensor:
         check_same_shape_as_values(self.weights, values, "weights")

src/torchjd/scalarization/_famo.py

@@ -128,7 +128,7 @@ def __init__(
                 f"Parameter `min_losses` should have shape {tuple(self._w.shape)} (matching the "
                 f"shape of the logits). Found `min_losses.shape = {tuple(min_losses.shape)}`."
             )
-        self.register_buffer("min_losses", min_losses)
+        self.register_buffer("min_losses", min_losses, persistent=False)
 
         self.lr = lr
         self.weight_decay = weight_decay

src/torchjd/scalarization/_pbi.py

@@ -51,14 +51,17 @@ class PBI(Scalarizer):
         :math:`10^{-6}` at that point and is negligible elsewhere.
     """
 
+    weights: Tensor
+    reference: Tensor | None
+
     def __init__(self, theta: float, weights: Tensor, reference: Tensor | None = None) -> None:
         if theta < 0.0:
             raise ValueError(f"Parameter `theta` should be non-negative. Found `theta = {theta}`.")
 
         super().__init__()
         self.theta = theta
-        self.weights = weights
-        self.reference = reference
+        self.register_buffer("weights", weights, persistent=False)
+        self.register_buffer("reference", reference, persistent=False)
 
     def forward(self, values: Tensor, /) -> Tensor:
         check_same_shape_as_values(self.weights, values, "weights")

src/torchjd/scalarization/_stch.py

@@ -38,6 +38,9 @@ class STCH(Scalarizer):
         is applied. If provided, it must have the same shape as the values passed at call time.
     """
 
+    weights: Tensor | None
+    reference: Tensor | None
+
     def __init__(
         self,
         mu: float,
@@ -49,8 +52,8 @@ def __init__(
 
         super().__init__()
         self.mu = mu
-        self.weights = weights
-        self.reference = reference
+        self.register_buffer("weights", weights, persistent=False)
+        self.register_buffer("reference", reference, persistent=False)
 
     def forward(self, values: Tensor, /) -> Tensor:
         check_same_shape_as_values(self.weights, values, "weights")

tests/unit/scalarization/test_constant.py

@@ -61,3 +61,11 @@ def test_representations() -> None:
 def test_str_with_non_vector_weights() -> None:
     assert str(Constant(weights=ones_((3, 4)))) == "Constant(weights of shape (3, 4))"
     assert str(Constant(weights=ones_(()))) == "Constant(weights of shape ())"
+
+
+def test_to_moves_weights() -> None:
+    s = Constant(weights=tensor_([1.0, 2.0], dtype=torch.float32)).to(torch.float64)
+
+    assert s.weights.dtype == torch.float64
+    assert any(buffer is s.weights for buffer in s.buffers())
+    assert len(s.state_dict()) == 0

tests/unit/scalarization/test_cosmos.py

@@ -86,3 +86,11 @@ def test_representations() -> None:
     s = COSMOS(lambda_=0.5, weights=torch.tensor([0.5, 0.5]))
     assert repr(s) == "COSMOS(lambda_=0.5, weights=tensor([0.5000, 0.5000]))"
     assert str(s) == "COSMOS"
+
+
+def test_to_moves_weights() -> None:
+    s = COSMOS(lambda_=0.5, weights=tensor_([0.5, 0.5], dtype=torch.float32)).to(torch.float64)
+
+    assert s.weights.dtype == torch.float64
+    assert any(buffer is s.weights for buffer in s.buffers())
+    assert len(s.state_dict()) == 0

tests/unit/scalarization/test_pbi.py

@@ -103,3 +103,19 @@ def test_representations() -> None:
     s = PBI(theta=5.0, weights=torch.tensor([0.5, 0.5]))
     assert repr(s) == "PBI(theta=5.0, weights=tensor([0.5000, 0.5000]), reference=None)"
     assert str(s) == "PBI"
+
+
+def test_to_moves_weights_and_reference() -> None:
+    s = PBI(
+        theta=5.0,
+        weights=tensor_([0.5, 0.5], dtype=torch.float32),
+        reference=tensor_([0.0, 1.0], dtype=torch.float32),
+    ).to(torch.float64)
+
+    assert s.reference is not None
+    assert s.weights.dtype == torch.float64
+    assert s.reference.dtype == torch.float64
+    buffers = list(s.buffers())
+    assert any(buffer is s.weights for buffer in buffers)
+    assert any(buffer is s.reference for buffer in buffers)
+    assert len(s.state_dict()) == 0

tests/unit/scalarization/test_stch.py

@@ -80,3 +80,20 @@ def test_representations() -> None:
     s = STCH(mu=0.5)
     assert repr(s) == "STCH(mu=0.5, weights=None, reference=None)"
     assert str(s) == "STCH"
+
+
+def test_to_moves_weights_and_reference() -> None:
+    s = STCH(
+        mu=1.0,
+        weights=tensor_([0.5, 0.5], dtype=torch.float32),
+        reference=tensor_([0.0, 1.0], dtype=torch.float32),
+    ).to(torch.float64)
+
+    assert s.weights is not None
+    assert s.reference is not None
+    assert s.weights.dtype == torch.float64
+    assert s.reference.dtype == torch.float64
+    buffers = list(s.buffers())
+    assert any(buffer is s.weights for buffer in buffers)
+    assert any(buffer is s.reference for buffer in buffers)
+    assert len(s.state_dict()) == 0