From 26f2e06576aad9d0c43cba2a79a325ec35fa50c7 Mon Sep 17 00:00:00 2001 From: ppraneth Date: Thu, 28 May 2026 19:41:37 +0530 Subject: [PATCH 1/7] add GLS support Signed-off-by: ppraneth --- CHANGELOG.md | 6 ++++ docs/source/docs/scalarization/gls.rst | 7 ++++ docs/source/docs/scalarization/index.rst | 1 + src/torchjd/scalarization/__init__.py | 9 ++---- src/torchjd/scalarization/_gls.py | 16 +++++++++ tests/unit/scalarization/test_gls.py | 41 ++++++++++++++++++++++++ 6 files changed, 73 insertions(+), 7 deletions(-) create mode 100644 docs/source/docs/scalarization/gls.rst create mode 100644 src/torchjd/scalarization/_gls.py create mode 100644 tests/unit/scalarization/test_gls.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 9204433a..03b9382d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,12 @@ changelog does not include internal changes that do not affect the user. ## [Unreleased] +### Added + +- Added `GLS` from [MultiNet++: Multi-Stream Feature Aggregation and Geometric Loss Strategy for + Multi-Task Learning](https://openaccess.thecvf.com/content_CVPRW_2019/papers/WAD/Chennupati_MultiNet_Multi-Stream_Feature_Aggregation_and_Geometric_Loss_Strategy_for_Multi-Task_CVPRW_2019_paper.pdf), + a `Scalarizer` that returns the geometric mean of the input tensor of values. + ## [0.12.0] - 2026-05-28 ### Added diff --git a/docs/source/docs/scalarization/gls.rst b/docs/source/docs/scalarization/gls.rst new file mode 100644 index 00000000..2ef36706 --- /dev/null +++ b/docs/source/docs/scalarization/gls.rst @@ -0,0 +1,7 @@ +:hide-toc: + +GLS +=== + +.. autoclass:: torchjd.scalarization.GLS + :members: __call__ diff --git a/docs/source/docs/scalarization/index.rst b/docs/source/docs/scalarization/index.rst index 11381bfa..0badb8d4 100644 --- a/docs/source/docs/scalarization/index.rst +++ b/docs/source/docs/scalarization/index.rst @@ -16,6 +16,7 @@ Abstract base class :maxdepth: 1 constant.rst + gls.rst mean.rst random.rst sum.rst diff --git a/src/torchjd/scalarization/__init__.py b/src/torchjd/scalarization/__init__.py index bf82aa11..645c95e9 100644 --- a/src/torchjd/scalarization/__init__.py +++ b/src/torchjd/scalarization/__init__.py @@ -20,15 +20,10 @@ """ from ._constant import Constant +from ._gls import GLS from ._mean import Mean from ._random import Random from ._scalarizer_base import Scalarizer from ._sum import Sum -__all__ = [ - "Constant", - "Mean", - "Random", - "Scalarizer", - "Sum", -] +__all__ = ["Constant", "GLS", "Mean", "Random", "Scalarizer", "Sum"] diff --git a/src/torchjd/scalarization/_gls.py b/src/torchjd/scalarization/_gls.py new file mode 100644 index 00000000..6d10bf94 --- /dev/null +++ b/src/torchjd/scalarization/_gls.py @@ -0,0 +1,16 @@ +import torch +from torch import Tensor + +from ._scalarizer_base import Scalarizer + + +class GLS(Scalarizer): + """ + :class:`~torchjd.scalarization.Scalarizer` that returns the geometric mean of the input tensor + of values, as defined in `MultiNet++: Multi-Stream Feature Aggregation and Geometric Loss + Strategy for Multi-Task Learning + `_. + """ + + def forward(self, values: Tensor, /) -> Tensor: + return torch.exp(torch.log(values).mean()) diff --git a/tests/unit/scalarization/test_gls.py b/tests/unit/scalarization/test_gls.py new file mode 100644 index 00000000..130ab949 --- /dev/null +++ b/tests/unit/scalarization/test_gls.py @@ -0,0 +1,41 @@ +import torch +from pytest import mark +from torch import Tensor +from utils.tensors import rand_, tensor_ + +from torchjd.scalarization import GLS + +from ._asserts import ( + assert_grad_flow, + assert_permutation_invariant, + assert_returns_scalar, +) + +shapes: list[list[int]] = [[], [5], [3, 4], [2, 3, 4]] +positive_inputs: list[Tensor] = [rand_(shape) + 1 for shape in shapes] + + +def test_value() -> None: + losses = tensor_([1.0, 2.0, 4.0]) + torch.testing.assert_close(GLS()(losses), tensor_(2.0)) + + +@mark.parametrize("losses", positive_inputs) +def test_expected_structure(losses: Tensor) -> None: + assert_returns_scalar(GLS(), losses) + + +@mark.parametrize("losses", positive_inputs) +def test_grad_flow(losses: Tensor) -> None: + assert_grad_flow(GLS(), losses) + + +@mark.parametrize("losses", positive_inputs) +def test_permutation_invariant(losses: Tensor) -> None: + assert_permutation_invariant(GLS(), losses) + + +def test_representations() -> None: + s = GLS() + assert repr(s) == "GLS()" + assert str(s) == "GLS" From ec4e3aaf1d2052e2114391aa01a91ddbcb1d425b Mon Sep 17 00:00:00 2001 From: ppraneth Date: Thu, 28 May 2026 19:52:07 +0530 Subject: [PATCH 2/7] add GLS support with edge cases Signed-off-by: ppraneth --- tests/unit/scalarization/test_gls.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/unit/scalarization/test_gls.py b/tests/unit/scalarization/test_gls.py index 130ab949..18c4c7df 100644 --- a/tests/unit/scalarization/test_gls.py +++ b/tests/unit/scalarization/test_gls.py @@ -35,6 +35,11 @@ def test_permutation_invariant(losses: Tensor) -> None: assert_permutation_invariant(GLS(), losses) +def test_propagates_nan_on_invalid_input() -> None: + # GLS is undefined for non-positive values; nan must propagate (no silent clamp). + assert GLS()(tensor_([1.0, -1.0])).isnan() + + def test_representations() -> None: s = GLS() assert repr(s) == "GLS()" From ac9f97aa6296456c0e22e2cc7432de6e14c79a0a Mon Sep 17 00:00:00 2001 From: ppraneth Date: Fri, 29 May 2026 06:55:06 +0530 Subject: [PATCH 3/7] minor fix Signed-off-by: ppraneth --- tests/unit/scalarization/test_gls.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/scalarization/test_gls.py b/tests/unit/scalarization/test_gls.py index 18c4c7df..0799b1c9 100644 --- a/tests/unit/scalarization/test_gls.py +++ b/tests/unit/scalarization/test_gls.py @@ -10,8 +10,8 @@ assert_permutation_invariant, assert_returns_scalar, ) +from ._inputs import shapes -shapes: list[list[int]] = [[], [5], [3, 4], [2, 3, 4]] positive_inputs: list[Tensor] = [rand_(shape) + 1 for shape in shapes] From 631609dd7d8844a2ba11009c1fee29c18b2f9853 Mon Sep 17 00:00:00 2001 From: Praneth Paruchuri Date: Fri, 29 May 2026 13:41:11 +0530 Subject: [PATCH 4/7] Update CHANGELOG.md Co-authored-by: Pierre Quinton --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 03b9382d..7db03865 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,7 @@ changelog does not include internal changes that do not affect the user. ### Added -- Added `GLS` from [MultiNet++: Multi-Stream Feature Aggregation and Geometric Loss Strategy for +- Added `GLS` studied in [MultiNet++: Multi-Stream Feature Aggregation and Geometric Loss Strategy for Multi-Task Learning](https://openaccess.thecvf.com/content_CVPRW_2019/papers/WAD/Chennupati_MultiNet_Multi-Stream_Feature_Aggregation_and_Geometric_Loss_Strategy_for_Multi-Task_CVPRW_2019_paper.pdf), a `Scalarizer` that returns the geometric mean of the input tensor of values. From 5966b0dfc04185ef9c6549c5d2edf7d6b7f9c77c Mon Sep 17 00:00:00 2001 From: Praneth Paruchuri Date: Fri, 29 May 2026 13:41:22 +0530 Subject: [PATCH 5/7] Update src/torchjd/scalarization/_gls.py Co-authored-by: Pierre Quinton --- src/torchjd/scalarization/_gls.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/torchjd/scalarization/_gls.py b/src/torchjd/scalarization/_gls.py index 6d10bf94..9f059f4e 100644 --- a/src/torchjd/scalarization/_gls.py +++ b/src/torchjd/scalarization/_gls.py @@ -7,7 +7,7 @@ class GLS(Scalarizer): """ :class:`~torchjd.scalarization.Scalarizer` that returns the geometric mean of the input tensor - of values, as defined in `MultiNet++: Multi-Stream Feature Aggregation and Geometric Loss + of values, as studied in `MultiNet++: Multi-Stream Feature Aggregation and Geometric Loss Strategy for Multi-Task Learning `_. """ From c21a57b90d4fc822ea39d5ce41661af60ef73e14 Mon Sep 17 00:00:00 2001 From: ppraneth Date: Fri, 29 May 2026 14:08:05 +0530 Subject: [PATCH 6/7] Rename GLS to GMean and update tests Signed-off-by: ppraneth --- CHANGELOG.md | 2 +- docs/source/docs/scalarization/gls.rst | 7 --- docs/source/docs/scalarization/gmean.rst | 7 +++ docs/source/docs/scalarization/index.rst | 2 +- src/torchjd/scalarization/__init__.py | 4 +- .../scalarization/{_gls.py => _gmean.py} | 4 +- tests/unit/scalarization/test_gls.py | 46 ---------------- tests/unit/scalarization/test_gmean.py | 53 +++++++++++++++++++ 8 files changed, 67 insertions(+), 58 deletions(-) delete mode 100644 docs/source/docs/scalarization/gls.rst create mode 100644 docs/source/docs/scalarization/gmean.rst rename src/torchjd/scalarization/{_gls.py => _gmean.py} (86%) delete mode 100644 tests/unit/scalarization/test_gls.py create mode 100644 tests/unit/scalarization/test_gmean.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 7db03865..12cae9c0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,7 @@ changelog does not include internal changes that do not affect the user. ### Added -- Added `GLS` studied in [MultiNet++: Multi-Stream Feature Aggregation and Geometric Loss Strategy for +- Added `GMean` studied in [MultiNet++: Multi-Stream Feature Aggregation and Geometric Loss Strategy for Multi-Task Learning](https://openaccess.thecvf.com/content_CVPRW_2019/papers/WAD/Chennupati_MultiNet_Multi-Stream_Feature_Aggregation_and_Geometric_Loss_Strategy_for_Multi-Task_CVPRW_2019_paper.pdf), a `Scalarizer` that returns the geometric mean of the input tensor of values. diff --git a/docs/source/docs/scalarization/gls.rst b/docs/source/docs/scalarization/gls.rst deleted file mode 100644 index 2ef36706..00000000 --- a/docs/source/docs/scalarization/gls.rst +++ /dev/null @@ -1,7 +0,0 @@ -:hide-toc: - -GLS -=== - -.. autoclass:: torchjd.scalarization.GLS - :members: __call__ diff --git a/docs/source/docs/scalarization/gmean.rst b/docs/source/docs/scalarization/gmean.rst new file mode 100644 index 00000000..6dc02150 --- /dev/null +++ b/docs/source/docs/scalarization/gmean.rst @@ -0,0 +1,7 @@ +:hide-toc: + +GMean +===== + +.. autoclass:: torchjd.scalarization.GMean + :members: __call__ diff --git a/docs/source/docs/scalarization/index.rst b/docs/source/docs/scalarization/index.rst index 0badb8d4..f3061a9a 100644 --- a/docs/source/docs/scalarization/index.rst +++ b/docs/source/docs/scalarization/index.rst @@ -16,7 +16,7 @@ Abstract base class :maxdepth: 1 constant.rst - gls.rst + gmean.rst mean.rst random.rst sum.rst diff --git a/src/torchjd/scalarization/__init__.py b/src/torchjd/scalarization/__init__.py index 645c95e9..bf5eed95 100644 --- a/src/torchjd/scalarization/__init__.py +++ b/src/torchjd/scalarization/__init__.py @@ -20,10 +20,10 @@ """ from ._constant import Constant -from ._gls import GLS +from ._gmean import GMean from ._mean import Mean from ._random import Random from ._scalarizer_base import Scalarizer from ._sum import Sum -__all__ = ["Constant", "GLS", "Mean", "Random", "Scalarizer", "Sum"] +__all__ = ["Constant", "GMean", "Mean", "Random", "Scalarizer", "Sum"] diff --git a/src/torchjd/scalarization/_gls.py b/src/torchjd/scalarization/_gmean.py similarity index 86% rename from src/torchjd/scalarization/_gls.py rename to src/torchjd/scalarization/_gmean.py index 9f059f4e..e9bb1508 100644 --- a/src/torchjd/scalarization/_gls.py +++ b/src/torchjd/scalarization/_gmean.py @@ -4,7 +4,7 @@ from ._scalarizer_base import Scalarizer -class GLS(Scalarizer): +class GMean(Scalarizer): """ :class:`~torchjd.scalarization.Scalarizer` that returns the geometric mean of the input tensor of values, as studied in `MultiNet++: Multi-Stream Feature Aggregation and Geometric Loss @@ -13,4 +13,6 @@ class GLS(Scalarizer): """ def forward(self, values: Tensor, /) -> Tensor: + if (values <= 0.0).any(): + return (values * 0.0).sum() return torch.exp(torch.log(values).mean()) diff --git a/tests/unit/scalarization/test_gls.py b/tests/unit/scalarization/test_gls.py deleted file mode 100644 index 0799b1c9..00000000 --- a/tests/unit/scalarization/test_gls.py +++ /dev/null @@ -1,46 +0,0 @@ -import torch -from pytest import mark -from torch import Tensor -from utils.tensors import rand_, tensor_ - -from torchjd.scalarization import GLS - -from ._asserts import ( - assert_grad_flow, - assert_permutation_invariant, - assert_returns_scalar, -) -from ._inputs import shapes - -positive_inputs: list[Tensor] = [rand_(shape) + 1 for shape in shapes] - - -def test_value() -> None: - losses = tensor_([1.0, 2.0, 4.0]) - torch.testing.assert_close(GLS()(losses), tensor_(2.0)) - - -@mark.parametrize("losses", positive_inputs) -def test_expected_structure(losses: Tensor) -> None: - assert_returns_scalar(GLS(), losses) - - -@mark.parametrize("losses", positive_inputs) -def test_grad_flow(losses: Tensor) -> None: - assert_grad_flow(GLS(), losses) - - -@mark.parametrize("losses", positive_inputs) -def test_permutation_invariant(losses: Tensor) -> None: - assert_permutation_invariant(GLS(), losses) - - -def test_propagates_nan_on_invalid_input() -> None: - # GLS is undefined for non-positive values; nan must propagate (no silent clamp). - assert GLS()(tensor_([1.0, -1.0])).isnan() - - -def test_representations() -> None: - s = GLS() - assert repr(s) == "GLS()" - assert str(s) == "GLS" diff --git a/tests/unit/scalarization/test_gmean.py b/tests/unit/scalarization/test_gmean.py new file mode 100644 index 00000000..aa096f6e --- /dev/null +++ b/tests/unit/scalarization/test_gmean.py @@ -0,0 +1,53 @@ +import torch +from pytest import mark +from torch import Tensor +from utils.tensors import rand_, tensor_ + +from torchjd.scalarization import GMean + +from ._asserts import ( + assert_grad_flow, + assert_permutation_invariant, + assert_returns_scalar, +) +from ._inputs import shapes + +positive_inputs: list[Tensor] = [rand_(shape) + 1 for shape in shapes] + + +def test_value() -> None: + losses = tensor_([1.0, 2.0, 4.0]) + torch.testing.assert_close(GMean()(losses), tensor_(2.0)) + + +@mark.parametrize("losses", positive_inputs) +def test_expected_structure(losses: Tensor) -> None: + assert_returns_scalar(GMean(), losses) + + +@mark.parametrize("losses", positive_inputs) +def test_grad_flow(losses: Tensor) -> None: + assert_grad_flow(GMean(), losses) + + +@mark.parametrize("losses", positive_inputs) +def test_permutation_invariant(losses: Tensor) -> None: + assert_permutation_invariant(GMean(), losses) + + +def test_returns_zero_on_non_positive_input() -> None: + # exp(log(x)) is undefined for x <= 0; the forward short-circuits to 0 in that case. + assert GMean()(tensor_([1.0, 0.0])) == 0.0 + assert GMean()(tensor_([1.0, -1.0])) == 0.0 + + +def test_grad_flow_on_zero_branch() -> None: + # The short-circuit branch must keep the autograd graph: backward() returns zero + # gradients rather than raising "does not have a grad_fn". + assert_grad_flow(GMean(), tensor_([1.0, 0.0, 2.0])) + + +def test_representations() -> None: + s = GMean() + assert repr(s) == "GMean()" + assert str(s) == "GMean" From 07861a37610f06891a27b5c75c46f3a0d2b852c1 Mon Sep 17 00:00:00 2001 From: ppraneth Date: Fri, 29 May 2026 16:05:02 +0530 Subject: [PATCH 7/7] Rename GMean to GeometricMean, update tests and docs Signed-off-by: ppraneth --- CHANGELOG.md | 5 +- .../docs/scalarization/geometric_mean.rst | 7 +++ docs/source/docs/scalarization/gmean.rst | 7 --- docs/source/docs/scalarization/index.rst | 2 +- src/torchjd/scalarization/__init__.py | 4 +- .../{_gmean.py => _geometric_mean.py} | 11 ++-- .../unit/scalarization/test_geometric_mean.py | 54 +++++++++++++++++++ tests/unit/scalarization/test_gmean.py | 53 ------------------ 8 files changed, 75 insertions(+), 68 deletions(-) create mode 100644 docs/source/docs/scalarization/geometric_mean.rst delete mode 100644 docs/source/docs/scalarization/gmean.rst rename src/torchjd/scalarization/{_gmean.py => _geometric_mean.py} (66%) create mode 100644 tests/unit/scalarization/test_geometric_mean.py delete mode 100644 tests/unit/scalarization/test_gmean.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 12cae9c0..ab5a536d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,8 +10,9 @@ changelog does not include internal changes that do not affect the user. ### Added -- Added `GMean` studied in [MultiNet++: Multi-Stream Feature Aggregation and Geometric Loss Strategy for - Multi-Task Learning](https://openaccess.thecvf.com/content_CVPRW_2019/papers/WAD/Chennupati_MultiNet_Multi-Stream_Feature_Aggregation_and_Geometric_Loss_Strategy_for_Multi-Task_CVPRW_2019_paper.pdf), +- Added `GeometricMean` (also known as GLS) studied in [MultiNet++: Multi-Stream Feature + Aggregation and Geometric Loss Strategy for Multi-Task + Learning](https://openaccess.thecvf.com/content_CVPRW_2019/papers/WAD/Chennupati_MultiNet_Multi-Stream_Feature_Aggregation_and_Geometric_Loss_Strategy_for_Multi-Task_CVPRW_2019_paper.pdf), a `Scalarizer` that returns the geometric mean of the input tensor of values. ## [0.12.0] - 2026-05-28 diff --git a/docs/source/docs/scalarization/geometric_mean.rst b/docs/source/docs/scalarization/geometric_mean.rst new file mode 100644 index 00000000..82622860 --- /dev/null +++ b/docs/source/docs/scalarization/geometric_mean.rst @@ -0,0 +1,7 @@ +:hide-toc: + +GeometricMean +============= + +.. autoclass:: torchjd.scalarization.GeometricMean + :members: __call__ diff --git a/docs/source/docs/scalarization/gmean.rst b/docs/source/docs/scalarization/gmean.rst deleted file mode 100644 index 6dc02150..00000000 --- a/docs/source/docs/scalarization/gmean.rst +++ /dev/null @@ -1,7 +0,0 @@ -:hide-toc: - -GMean -===== - -.. autoclass:: torchjd.scalarization.GMean - :members: __call__ diff --git a/docs/source/docs/scalarization/index.rst b/docs/source/docs/scalarization/index.rst index f3061a9a..8fd87dc8 100644 --- a/docs/source/docs/scalarization/index.rst +++ b/docs/source/docs/scalarization/index.rst @@ -16,7 +16,7 @@ Abstract base class :maxdepth: 1 constant.rst - gmean.rst + geometric_mean.rst mean.rst random.rst sum.rst diff --git a/src/torchjd/scalarization/__init__.py b/src/torchjd/scalarization/__init__.py index bf5eed95..337d38ca 100644 --- a/src/torchjd/scalarization/__init__.py +++ b/src/torchjd/scalarization/__init__.py @@ -20,10 +20,10 @@ """ from ._constant import Constant -from ._gmean import GMean +from ._geometric_mean import GeometricMean from ._mean import Mean from ._random import Random from ._scalarizer_base import Scalarizer from ._sum import Sum -__all__ = ["Constant", "GMean", "Mean", "Random", "Scalarizer", "Sum"] +__all__ = ["Constant", "GeometricMean", "Mean", "Random", "Scalarizer", "Sum"] diff --git a/src/torchjd/scalarization/_gmean.py b/src/torchjd/scalarization/_geometric_mean.py similarity index 66% rename from src/torchjd/scalarization/_gmean.py rename to src/torchjd/scalarization/_geometric_mean.py index e9bb1508..f54d9fca 100644 --- a/src/torchjd/scalarization/_gmean.py +++ b/src/torchjd/scalarization/_geometric_mean.py @@ -4,15 +4,20 @@ from ._scalarizer_base import Scalarizer -class GMean(Scalarizer): +class GeometricMean(Scalarizer): """ :class:`~torchjd.scalarization.Scalarizer` that returns the geometric mean of the input tensor of values, as studied in `MultiNet++: Multi-Stream Feature Aggregation and Geometric Loss Strategy for Multi-Task Learning `_. + + This method is also known as GLS (Geometric Loss Strategy). """ def forward(self, values: Tensor, /) -> Tensor: - if (values <= 0.0).any(): - return (values * 0.0).sum() + if (values < 1e-12).any(): + raise ValueError( + "GeometricMean is only defined for strictly positive values. Found a value " + "below 1e-12 in the input." + ) return torch.exp(torch.log(values).mean()) diff --git a/tests/unit/scalarization/test_geometric_mean.py b/tests/unit/scalarization/test_geometric_mean.py new file mode 100644 index 00000000..1fb99ab9 --- /dev/null +++ b/tests/unit/scalarization/test_geometric_mean.py @@ -0,0 +1,54 @@ +import torch +from pytest import mark, raises +from torch import Tensor +from utils.tensors import rand_, tensor_ + +from torchjd.scalarization import GeometricMean + +from ._asserts import ( + assert_grad_flow, + assert_permutation_invariant, + assert_returns_scalar, +) +from ._inputs import shapes + +positive_inputs: list[Tensor] = [rand_(shape) + 1 for shape in shapes] + + +def test_value() -> None: + losses = tensor_([1.0, 2.0, 4.0]) + torch.testing.assert_close(GeometricMean()(losses), tensor_(2.0)) + + +@mark.parametrize("losses", positive_inputs) +def test_expected_structure(losses: Tensor) -> None: + assert_returns_scalar(GeometricMean(), losses) + + +@mark.parametrize("losses", positive_inputs) +def test_grad_flow(losses: Tensor) -> None: + assert_grad_flow(GeometricMean(), losses) + + +@mark.parametrize("losses", positive_inputs) +def test_permutation_invariant(losses: Tensor) -> None: + assert_permutation_invariant(GeometricMean(), losses) + + +@mark.parametrize( + "invalid", + [ + tensor_([1.0, 0.0]), + tensor_([1.0, -1.0]), + tensor_([1.0, 1e-13]), + ], +) +def test_raises_on_non_positive_input(invalid: Tensor) -> None: + with raises(ValueError): + GeometricMean()(invalid) + + +def test_representations() -> None: + s = GeometricMean() + assert repr(s) == "GeometricMean()" + assert str(s) == "GeometricMean" diff --git a/tests/unit/scalarization/test_gmean.py b/tests/unit/scalarization/test_gmean.py deleted file mode 100644 index aa096f6e..00000000 --- a/tests/unit/scalarization/test_gmean.py +++ /dev/null @@ -1,53 +0,0 @@ -import torch -from pytest import mark -from torch import Tensor -from utils.tensors import rand_, tensor_ - -from torchjd.scalarization import GMean - -from ._asserts import ( - assert_grad_flow, - assert_permutation_invariant, - assert_returns_scalar, -) -from ._inputs import shapes - -positive_inputs: list[Tensor] = [rand_(shape) + 1 for shape in shapes] - - -def test_value() -> None: - losses = tensor_([1.0, 2.0, 4.0]) - torch.testing.assert_close(GMean()(losses), tensor_(2.0)) - - -@mark.parametrize("losses", positive_inputs) -def test_expected_structure(losses: Tensor) -> None: - assert_returns_scalar(GMean(), losses) - - -@mark.parametrize("losses", positive_inputs) -def test_grad_flow(losses: Tensor) -> None: - assert_grad_flow(GMean(), losses) - - -@mark.parametrize("losses", positive_inputs) -def test_permutation_invariant(losses: Tensor) -> None: - assert_permutation_invariant(GMean(), losses) - - -def test_returns_zero_on_non_positive_input() -> None: - # exp(log(x)) is undefined for x <= 0; the forward short-circuits to 0 in that case. - assert GMean()(tensor_([1.0, 0.0])) == 0.0 - assert GMean()(tensor_([1.0, -1.0])) == 0.0 - - -def test_grad_flow_on_zero_branch() -> None: - # The short-circuit branch must keep the autograd graph: backward() returns zero - # gradients rather than raising "does not have a grad_fn". - assert_grad_flow(GMean(), tensor_([1.0, 0.0, 2.0])) - - -def test_representations() -> None: - s = GMean() - assert repr(s) == "GMean()" - assert str(s) == "GMean"