Make _pref_vector_to_weighting take a default weighting as param (#229)

Author: ValerianRey

src/torchjd/aggregation/_pref_vector_utils.py

@@ -1,7 +1,7 @@
 from torch import Tensor
 
+from .bases import _Weighting
 from .constant import _ConstantWeighting
-from .mean import _MeanWeighting
 
 
 def _check_pref_vector(pref_vector: Tensor | None) -> None:
@@ -15,12 +15,13 @@ def _check_pref_vector(pref_vector: Tensor | None) -> None:
             )
 
 
-def _pref_vector_to_weighting(pref_vector: Tensor | None) -> _ConstantWeighting | _MeanWeighting:
-    """Returns the weighting associated to a given preference vector."""
+def _pref_vector_to_weighting(pref_vector: Tensor | None, default: _Weighting) -> _Weighting:
+    """
+    Returns the weighting associated to a given preference vector, with a fallback to a default
+    weighting if the preference vector is None.
+    """
 
     if pref_vector is None:
-        weighting = _MeanWeighting()
+        return default
     else:
-        weighting = _ConstantWeighting(pref_vector)
-
-    return weighting
+        return _ConstantWeighting(pref_vector)

src/torchjd/aggregation/aligned_mtl.py

@@ -32,6 +32,7 @@
 from ._pref_vector_utils import _check_pref_vector, _pref_vector_to_weighting
 from ._str_utils import _vector_to_str
 from .bases import _WeightedAggregator, _Weighting
+from .mean import _MeanWeighting
 
 
 class AlignedMTL(_WeightedAggregator):
@@ -63,7 +64,7 @@ class AlignedMTL(_WeightedAggregator):
 
     def __init__(self, pref_vector: Tensor | None = None):
         _check_pref_vector(pref_vector)
-        weighting = _pref_vector_to_weighting(pref_vector)
+        weighting = _pref_vector_to_weighting(pref_vector, default=_MeanWeighting())
         self._pref_vector = pref_vector
 
         super().__init__(weighting=_AlignedMTLWrapper(weighting))

src/torchjd/aggregation/dualproj.py

@@ -9,6 +9,7 @@
 from ._pref_vector_utils import _check_pref_vector, _pref_vector_to_weighting
 from ._str_utils import _vector_to_str
 from .bases import _WeightedAggregator, _Weighting
+from .mean import _MeanWeighting
 
 
 class DualProj(_WeightedAggregator):
@@ -50,7 +51,7 @@ def __init__(
         solver: Literal["quadprog"] = "quadprog",
     ):
         _check_pref_vector(pref_vector)
-        weighting = _pref_vector_to_weighting(pref_vector)
+        weighting = _pref_vector_to_weighting(pref_vector, default=_MeanWeighting())
         self._pref_vector = pref_vector
 
         super().__init__(

src/torchjd/aggregation/upgrad.py

@@ -9,6 +9,7 @@
 from ._pref_vector_utils import _check_pref_vector, _pref_vector_to_weighting
 from ._str_utils import _vector_to_str
 from .bases import _WeightedAggregator, _Weighting
+from .mean import _MeanWeighting
 
 
 class UPGrad(_WeightedAggregator):
@@ -49,7 +50,7 @@ def __init__(
         solver: Literal["quadprog"] = "quadprog",
     ):
         _check_pref_vector(pref_vector)
-        weighting = _pref_vector_to_weighting(pref_vector)
+        weighting = _pref_vector_to_weighting(pref_vector, default=_MeanWeighting())
         self._pref_vector = pref_vector
 
         super().__init__(