From 9a802ccb409cedb57fc47ab09cc58e85a5574638 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Val=C3=A9rian=20Rey?= Date: Wed, 7 Oct 2026 17:27:31 +0200 Subject: [PATCH] Extract shape check in scalarizers COSMOS, PBI and STCH each repeated the same check that `weights` or `reference` has the shape of `values`. It is now done by `check_same_shape_as_values` in `scalarization/_utils.py`. The error messages are unchanged. Co-Authored-By: Claude Opus 5.5 --- src/torchjd/scalarization/_cosmos.py | 8 ++------ src/torchjd/scalarization/_pbi.py | 15 +++------------ src/torchjd/scalarization/_stch.py | 15 +++------------ src/torchjd/scalarization/_utils.py | 18 ++++++++++++++++++ 4 files changed, 26 insertions(+), 30 deletions(-) create mode 100644 src/torchjd/scalarization/_utils.py diff --git a/src/torchjd/scalarization/_cosmos.py b/src/torchjd/scalarization/_cosmos.py index be7bc3dc..2de5077c 100644 --- a/src/torchjd/scalarization/_cosmos.py +++ b/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) diff --git a/src/torchjd/scalarization/_pbi.py b/src/torchjd/scalarization/_pbi.py index defa5796..fe8f5025 100644 --- a/src/torchjd/scalarization/_pbi.py +++ b/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() diff --git a/src/torchjd/scalarization/_stch.py b/src/torchjd/scalarization/_stch.py index 9df18ee1..25ec95bc 100644 --- a/src/torchjd/scalarization/_stch.py +++ b/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()) diff --git a/src/torchjd/scalarization/_utils.py b/src/torchjd/scalarization/_utils.py new file mode 100644 index 00000000..7491413e --- /dev/null +++ b/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)}`." + )