Fix(aggregation): Fix some aggregator __repr__ (#362)

* Fix NashMTL.__repr__ to also include max_norm, update_weights_every and optim_niter. Adapt its representation test.
* Fix GradDrop.__repr__ to also include f. Adapt its representation test using a regex.
* Add _epsilon to MGDA
* Add _weights to Constant

Author: ValerianRey

src/torchjd/aggregation/_constant.py

@@ -29,12 +29,13 @@ class Constant(WeightedAggregator):
 
     def __init__(self, weights: Tensor):
         super().__init__(weighting=_ConstantWeighting(weights=weights))
+        self._weights = weights
 
     def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(weights={repr(self.weighting.weights)})"
+        return f"{self.__class__.__name__}(weights={repr(self._weights)})"
 
     def __str__(self) -> str:
-        weights_str = vector_to_str(self.weighting.weights)
+        weights_str = vector_to_str(self._weights)
         return f"{self.__class__.__name__}([{weights_str}])"
 
 

src/torchjd/aggregation/_graddrop.py

@@ -82,7 +82,7 @@ def _check_matrix_has_enough_rows(self, matrix: Tensor) -> None:
             )
 
     def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(leak={repr(self.leak)})"
+        return f"{self.__class__.__name__}(f={repr(self.f)}, leak={repr(self.leak)})"
 
     def __str__(self) -> str:
         if self.leak is None:

src/torchjd/aggregation/_mgda.py

@@ -33,13 +33,11 @@ class MGDA(GramianWeightedAggregator):
 
     def __init__(self, epsilon: float = 0.001, max_iters: int = 100):
         super().__init__(_MGDAWeighting(epsilon=epsilon, max_iters=max_iters))
+        self._epsilon = epsilon
         self._max_iters = max_iters
 
     def __repr__(self) -> str:
-        return (
-            f"{self.__class__.__name__}(epsilon={self.weighting.weighting.epsilon}, max_iters="
-            f"{self._max_iters})"
-        )
+        return f"{self.__class__.__name__}(epsilon={self._epsilon}, max_iters={self._max_iters})"
 
 
 class _MGDAWeighting(Weighting[PSDMatrix]):

src/torchjd/aggregation/_nash_mtl.py

@@ -96,6 +96,10 @@ def __init__(
                 optim_niter=optim_niter,
             )
         )
+        self._n_tasks = n_tasks
+        self._max_norm = max_norm
+        self._update_weights_every = update_weights_every
+        self._optim_niter = optim_niter
 
         # This prevents considering the computed weights as constant w.r.t. the matrix.
         self.register_full_backward_pre_hook(raise_non_differentiable_error)
@@ -105,7 +109,10 @@ def reset(self) -> None:
         self.weighting.reset()
 
     def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(n_tasks={self.weighting.n_tasks})"
+        return (
+            f"{self.__class__.__name__}(n_tasks={self._n_tasks}, max_norm={self._max_norm}, "
+            f"update_weights_every={self._update_weights_every}, optim_niter={self._optim_niter})"
+        )
 
 
 class _NashMTLWeighting(Weighting[Matrix]):

tests/unit/aggregation/test_graddrop.py

@@ -1,3 +1,4 @@
+import re
 from contextlib import nullcontext as does_not_raise
 
 import torch
@@ -69,9 +70,13 @@ def test_matrix_shape_check(leak_shape: list[int], n_rows: int, expectation: Exc
 
 def test_representations():
     A = GradDrop(leak=torch.tensor([0.0, 1.0], device="cpu"))
-    assert repr(A) == "GradDrop(leak=tensor([0., 1.]))"
+    assert re.match(
+        r"GradDrop\(f=<function _identity at 0x[0-9a-fA-F]+>, leak=tensor\(\[0\., 1\.\]\)\)",
+        repr(A),
+    )
+
     assert str(A) == "GradDrop([0., 1.])"
 
     A = GradDrop()
-    assert repr(A) == "GradDrop(leak=None)"
+    assert re.match(r"GradDrop\(f=<function _identity at 0x[0-9a-fA-F]+>, leak=None\)", repr(A))
     assert str(A) == "GradDrop"

tests/unit/aggregation/test_nash_mtl.py

@@ -55,6 +55,6 @@ def test_nash_mtl_reset():
 
 
 def test_representations():
-    A = NashMTL(n_tasks=2)
-    assert repr(A) == "NashMTL(n_tasks=2)"
+    A = NashMTL(n_tasks=2, max_norm=1.5, update_weights_every=2, optim_niter=5)
+    assert repr(A) == "NashMTL(n_tasks=2, max_norm=1.5, update_weights_every=2, optim_niter=5)"
     assert str(A) == "NashMTL"