diff --git a/src/torchjd/aggregation/_cagrad.py b/src/torchjd/aggregation/_cagrad.py index 40a5d5b0..150ebb76 100644 --- a/src/torchjd/aggregation/_cagrad.py +++ b/src/torchjd/aggregation/_cagrad.py @@ -135,5 +135,5 @@ def __repr__(self) -> str: return f"{self.__class__.__name__}(c={self.c}, norm_eps={self.norm_eps})" def __str__(self) -> str: - c_str = str(self.c).rstrip("0") + c_str = str(self.c).removesuffix(".0") return f"CAGrad{c_str}" diff --git a/tests/unit/aggregation/test_cagrad.py b/tests/unit/aggregation/test_cagrad.py index 6cb5e968..921cce13 100644 --- a/tests/unit/aggregation/test_cagrad.py +++ b/tests/unit/aggregation/test_cagrad.py @@ -60,6 +60,22 @@ def test_representations() -> None: assert str(A) == "CAGrad0.5" +@mark.parametrize( + ["c", "expected"], + [ + (0, "CAGrad0"), + (0.0, "CAGrad0"), + (0.5, "CAGrad0.5"), + (1, "CAGrad1"), + (1.0, "CAGrad1"), + (10, "CAGrad10"), + (100.0, "CAGrad100"), + ], +) +def test_str_keeps_significant_zeros(c: float, expected: str) -> None: + assert str(CAGrad(c=c)) == expected + + def test_c_setter_updates_value() -> None: A = CAGrad(c=0.5) A.c = 1.25