From 1f6b634ff3dff0d3ad47112df711af1592d55017 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Val=C3=A9rian=20Rey?= Date: Wed, 7 Oct 2026 17:30:26 +0200 Subject: [PATCH] Rename simplex projection helper to project_onto_simplex `_projection2simplex` is imported by MoDo, SDMGrad and the tests, and its sibling helper in `aggregation/_utils` (`pref_vector_to_weighting`) has no leading underscore and spells out "to". Also use `clamp` instead of taking the maximum with a new tensor of zeros. Outputs are unchanged. Co-Authored-By: Claude Opus 5.5 --- src/torchjd/aggregation/_modo.py | 4 ++-- src/torchjd/aggregation/_sdmgrad.py | 4 ++-- src/torchjd/aggregation/_utils/simplex.py | 4 ++-- tests/unit/aggregation/_utils/test_simplex.py | 8 ++++---- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/src/torchjd/aggregation/_modo.py b/src/torchjd/aggregation/_modo.py index 84d64e10..0d4c76bc 100644 --- a/src/torchjd/aggregation/_modo.py +++ b/src/torchjd/aggregation/_modo.py @@ -11,7 +11,7 @@ from torchjd.aggregation._mixins import _NonDifferentiable from torchjd.linalg import Matrix -from ._utils.simplex import _projection2simplex +from ._utils.simplex import project_onto_simplex from ._weighting_bases import _MatrixWeighting @@ -167,7 +167,7 @@ def forward(self, matrix: Matrix, /) -> Tensor: lambd = cast(Tensor, self._lambda) grad = matrix @ lambd + self._rho * lambd - lambd = _projection2simplex(lambd - self._gamma * grad) + lambd = project_onto_simplex(lambd - self._gamma * grad) self._lambda = lambd return lambd diff --git a/src/torchjd/aggregation/_sdmgrad.py b/src/torchjd/aggregation/_sdmgrad.py index 0431f296..15120c3a 100644 --- a/src/torchjd/aggregation/_sdmgrad.py +++ b/src/torchjd/aggregation/_sdmgrad.py @@ -11,7 +11,7 @@ from torchjd.aggregation._mixins import _NonDifferentiable from torchjd.linalg import Matrix -from ._utils.simplex import _projection2simplex +from ._utils.simplex import project_onto_simplex from ._weighting_bases import _MatrixWeighting @@ -184,7 +184,7 @@ def forward(self, matrix: Matrix, /) -> Tensor: for _ in range(self._n_iter): grad = a @ (w + self._lambda * w_tilde) velocity = grad if velocity is None else self._momentum * velocity + grad - w = _projection2simplex(w - self._lr * velocity) + w = project_onto_simplex(w - self._lr * velocity) self._w = w return (w + self._lambda * w_tilde) / (1.0 + self._lambda) diff --git a/src/torchjd/aggregation/_utils/simplex.py b/src/torchjd/aggregation/_utils/simplex.py index b146d307..3a89dc83 100644 --- a/src/torchjd/aggregation/_utils/simplex.py +++ b/src/torchjd/aggregation/_utils/simplex.py @@ -2,7 +2,7 @@ from torch import Tensor -def _projection2simplex(y: Tensor) -> Tensor: +def project_onto_simplex(y: Tensor) -> Tensor: """Euclidean projection of ``y`` onto the probability simplex.""" m = len(y) @@ -15,4 +15,4 @@ def _projection2simplex(y: Tensor) -> Tensor: if tmax > sorted_y[i + 1]: tmax_f = tmax break - return torch.max(y - tmax_f, y.new_zeros(m)) + return (y - tmax_f).clamp(min=0.0) diff --git a/tests/unit/aggregation/_utils/test_simplex.py b/tests/unit/aggregation/_utils/test_simplex.py index 7f3907f5..78b5f2f1 100644 --- a/tests/unit/aggregation/_utils/test_simplex.py +++ b/tests/unit/aggregation/_utils/test_simplex.py @@ -1,19 +1,19 @@ from torch.testing import assert_close from utils.tensors import tensor_ -from torchjd.aggregation._utils.simplex import _projection2simplex +from torchjd.aggregation._utils.simplex import project_onto_simplex -def test_projection2simplex_known_values() -> None: +def test_project_onto_simplex_known_values() -> None: """The simplex projection matches hand-computed Euclidean projections.""" # Already-positive input: the deficit (1 - sum) is spread equally, no clamping. assert_close( - _projection2simplex(tensor_([0.5, 0.1, 0.1])), + project_onto_simplex(tensor_([0.5, 0.1, 0.1])), tensor_([0.6, 0.2, 0.2]), ) # Input with a negative entry: it gets clamped to zero. assert_close( - _projection2simplex(tensor_([1.0, 0.0, -0.5])), + project_onto_simplex(tensor_([1.0, 0.0, -0.5])), tensor_([1.0, 0.0, 0.0]), )