refactor(scalarization): Extract shape check in scalarizers

src/torchjd/scalarization/_cosmos.py

@@ -2,6 +2,7 @@
 from torch.nn.functional import cosine_similarity
 
 from ._scalarizer_base import Scalarizer
+from ._utils import check_same_shape_as_values
 
 
 class COSMOS(Scalarizer):
@@ -52,12 +53,7 @@ def __init__(self, lambda_: float, weights: Tensor) -> None:
         self.weights = weights
 
     def forward(self, values: Tensor, /) -> Tensor:
-        if self.weights.shape != values.shape:
-            raise ValueError(
-                f"Parameter `weights` should have the same shape as `values`. Found "
-                f"`weights.shape = {tuple(self.weights.shape)}` and `values.shape = "
-                f"{tuple(values.shape)}`."
-            )
+        check_same_shape_as_values(self.weights, values, "weights")
 
         weighted_sum = (self.weights * values).sum()
         cosine = cosine_similarity(self.weights.flatten(), values.flatten(), dim=0)

src/torchjd/scalarization/_pbi.py

@@ -2,6 +2,7 @@
 from torch import Tensor
 
 from ._scalarizer_base import Scalarizer
+from ._utils import check_same_shape_as_values
 
 _EPSILON = 1e-12
 
@@ -60,18 +61,8 @@ def __init__(self, theta: float, weights: Tensor, reference: Tensor | None = Non
         self.reference = reference
 
     def forward(self, values: Tensor, /) -> Tensor:
-        if self.weights.shape != values.shape:
-            raise ValueError(
-                f"Parameter `weights` should have the same shape as `values`. Found "
-                f"`weights.shape = {tuple(self.weights.shape)}` and `values.shape = "
-                f"{tuple(values.shape)}`."
-            )
-        if self.reference is not None and self.reference.shape != values.shape:
-            raise ValueError(
-                f"Parameter `reference` should have the same shape as `values`. Found "
-                f"`reference.shape = {tuple(self.reference.shape)}` and `values.shape = "
-                f"{tuple(values.shape)}`."
-            )
+        check_same_shape_as_values(self.weights, values, "weights")
+        check_same_shape_as_values(self.reference, values, "reference")
 
         shifted = values if self.reference is None else values - self.reference
         f = shifted.flatten()

src/torchjd/scalarization/_stch.py

@@ -2,6 +2,7 @@
 from torch import Tensor
 
 from ._scalarizer_base import Scalarizer
+from ._utils import check_same_shape_as_values
 
 
 class STCH(Scalarizer):
@@ -52,18 +53,8 @@ def __init__(
         self.reference = reference
 
     def forward(self, values: Tensor, /) -> Tensor:
-        if self.weights is not None and self.weights.shape != values.shape:
-            raise ValueError(
-                f"Parameter `weights` should have the same shape as `values`. Found "
-                f"`weights.shape = {tuple(self.weights.shape)}` and `values.shape = "
-                f"{tuple(values.shape)}`."
-            )
-        if self.reference is not None and self.reference.shape != values.shape:
-            raise ValueError(
-                f"Parameter `reference` should have the same shape as `values`. Found "
-                f"`reference.shape = {tuple(self.reference.shape)}` and `values.shape = "
-                f"{tuple(values.shape)}`."
-            )
+        check_same_shape_as_values(self.weights, values, "weights")
+        check_same_shape_as_values(self.reference, values, "reference")
 
         if self.weights is None:
             weights = torch.full_like(values, 1.0 / values.numel())

src/torchjd/scalarization/_utils.py

@@ -0,0 +1,18 @@
+from torch import Tensor
+
+
+def check_same_shape_as_values(tensor: Tensor | None, values: Tensor, name: str) -> None:
+    """
+    Checks that ``tensor`` has the same shape as ``values``. Nothing is checked if ``tensor`` is
+    ``None``.
+
+    :param tensor: The tensor to check.
+    :param values: The tensor of values to scalarize.
+    :param name: The name of the parameter corresponding to ``tensor``, used in the error message.
+    """
+
+    if tensor is not None and tensor.shape != values.shape:
+        raise ValueError(
+            f"Parameter `{name}` should have the same shape as `values`. Found "
+            f"`{name}.shape = {tuple(tensor.shape)}` and `values.shape = {tuple(values.shape)}`."
+        )