refactor(aggregation): Move utils of aggregation to _utils package (#338)

PierreQuinton · ValerianRey · web-flow · commit 24a5c5470534 · 2025-05-10T01:47:50.000+02:00
* Create `_utils` package in `aggregation`.
* Move utilitary files in `_utils` and remove their `_utils` suffix.
* Create the `_utils` package in `aggregation` in the unit tests, and move tests there.
---------
Co-authored-by: Valérian Rey &lt;valerian.rey@gmail.com&gt;
diff --git a/src/torchjd/aggregation/__init__.py b/src/torchjd/aggregation/__init__.py
@@ -1,4 +1,6 @@
-from ._check_dependencies import _OptionalDepsNotInstalledError
+from ._utils.check_dependencies import (
+    OptionalDepsNotInstalledError as _OptionalDepsNotInstalledError,
+)
 from .aligned_mtl import AlignedMTL
 from .bases import Aggregator
 from .config import ConFIG
diff --git a/src/torchjd/aggregation/_utils/__init__.py b/src/torchjd/aggregation/_utils/__init__.py
diff --git a/src/torchjd/aggregation/_utils/check_dependencies.py b/src/torchjd/aggregation/_utils/check_dependencies.py
@@ -1,7 +1,7 @@
 from importlib.util import find_spec
 
 
-class _OptionalDepsNotInstalledError(ModuleNotFoundError):
+class OptionalDepsNotInstalledError(ModuleNotFoundError):
     pass
 
 
@@ -13,4 +13,4 @@ def check_dependencies_are_installed(dependency_names: list[str]) -> None:
     """
 
     if any(find_spec(name) is None for name in dependency_names):
-        raise _OptionalDepsNotInstalledError()
+        raise OptionalDepsNotInstalledError()
diff --git a/src/torchjd/aggregation/_utils/dual_cone.py b/src/torchjd/aggregation/_utils/dual_cone.py
diff --git a/src/torchjd/aggregation/_utils/gramian.py b/src/torchjd/aggregation/_utils/gramian.py
diff --git a/src/torchjd/aggregation/_utils/non_differentiable.py b/src/torchjd/aggregation/_utils/non_differentiable.py
diff --git a/src/torchjd/aggregation/_utils/pref_vector.py b/src/torchjd/aggregation/_utils/pref_vector.py
@@ -1,8 +1,9 @@
 from torch import Tensor
 
-from ._str_utils import _vector_to_str
-from .bases import _Weighting
-from .constant import _ConstantWeighting
+from torchjd.aggregation.bases import _Weighting
+from torchjd.aggregation.constant import _ConstantWeighting
+
+from .str import vector_to_str
 
 
 def pref_vector_to_weighting(pref_vector: Tensor | None, default: _Weighting) -> _Weighting:
@@ -28,4 +29,4 @@ def pref_vector_to_str_suffix(pref_vector: Tensor | None) -> str:
     if pref_vector is None:
         return ""
     else:
-        return f"([{_vector_to_str(pref_vector)}])"
+        return f"([{vector_to_str(pref_vector)}])"
diff --git a/src/torchjd/aggregation/_utils/str.py b/src/torchjd/aggregation/_utils/str.py
@@ -1,7 +1,7 @@
 from torch import Tensor
 
 
-def _vector_to_str(vector: Tensor) -> str:
+def vector_to_str(vector: Tensor) -> str:
     """
     Transforms a Tensor of the form `tensor([1.23456, 1.0, ...])` into a string of the form
     `1.23, 1., ...`
diff --git a/src/torchjd/aggregation/aligned_mtl.py b/src/torchjd/aggregation/aligned_mtl.py
@@ -28,8 +28,8 @@
 import torch
 from torch import Tensor
 
-from ._gramian_utils import compute_gramian
-from ._pref_vector_utils import pref_vector_to_str_suffix, pref_vector_to_weighting
+from ._utils.gramian import compute_gramian
+from ._utils.pref_vector import pref_vector_to_str_suffix, pref_vector_to_weighting
 from .bases import _WeightedAggregator, _Weighting
 from .mean import _MeanWeighting
 
diff --git a/src/torchjd/aggregation/cagrad.py b/src/torchjd/aggregation/cagrad.py
@@ -1,4 +1,4 @@
-from ._check_dependencies import check_dependencies_are_installed  # noqa
+from ._utils.check_dependencies import check_dependencies_are_installed
 
 check_dependencies_are_installed(["cvxpy", "clarabel"])
 
@@ -7,8 +7,8 @@
 import torch
 from torch import Tensor
 
-from ._gramian_utils import compute_gramian, normalize
-from ._non_differentiable import raise_non_differentiable_error
+from ._utils.gramian import compute_gramian, normalize
+from ._utils.non_differentiable import raise_non_differentiable_error
 from .bases import _WeightedAggregator, _Weighting
 
 
diff --git a/src/torchjd/aggregation/config.py b/src/torchjd/aggregation/config.py
@@ -28,8 +28,8 @@
 import torch
 from torch import Tensor
 
-from ._non_differentiable import raise_non_differentiable_error
-from ._pref_vector_utils import pref_vector_to_str_suffix, pref_vector_to_weighting
+from ._utils.non_differentiable import raise_non_differentiable_error
+from ._utils.pref_vector import pref_vector_to_str_suffix, pref_vector_to_weighting
 from .bases import Aggregator
 from .sum import _SumWeighting
 
diff --git a/src/torchjd/aggregation/constant.py b/src/torchjd/aggregation/constant.py
@@ -1,6 +1,6 @@
 from torch import Tensor
 
-from ._str_utils import _vector_to_str
+from ._utils.str import vector_to_str
 from .bases import _WeightedAggregator, _Weighting
 
 
@@ -33,7 +33,7 @@ def __repr__(self) -> str:
         return f"{self.__class__.__name__}(weights={repr(self.weighting.weights)})"
 
     def __str__(self) -> str:
-        weights_str = _vector_to_str(self.weighting.weights)
+        weights_str = vector_to_str(self.weighting.weights)
         return f"{self.__class__.__name__}([{weights_str}])"
 
 
diff --git a/src/torchjd/aggregation/dualproj.py b/src/torchjd/aggregation/dualproj.py
@@ -2,10 +2,10 @@
 
 from torch import Tensor
 
-from ._dual_cone_utils import project_weights
-from ._gramian_utils import compute_gramian, normalize, regularize
-from ._non_differentiable import raise_non_differentiable_error
-from ._pref_vector_utils import pref_vector_to_str_suffix, pref_vector_to_weighting
+from ._utils.dual_cone import project_weights
+from ._utils.gramian import compute_gramian, normalize, regularize
+from ._utils.non_differentiable import raise_non_differentiable_error
+from ._utils.pref_vector import pref_vector_to_str_suffix, pref_vector_to_weighting
 from .bases import _WeightedAggregator, _Weighting
 from .mean import _MeanWeighting
 
diff --git a/src/torchjd/aggregation/graddrop.py b/src/torchjd/aggregation/graddrop.py
@@ -3,7 +3,7 @@
 import torch
 from torch import Tensor
 
-from ._non_differentiable import raise_non_differentiable_error
+from ._utils.non_differentiable import raise_non_differentiable_error
 from .bases import Aggregator
 
 
diff --git a/src/torchjd/aggregation/imtl_g.py b/src/torchjd/aggregation/imtl_g.py
@@ -1,8 +1,8 @@
 import torch
 from torch import Tensor
 
-from ._gramian_utils import compute_gramian
-from ._non_differentiable import raise_non_differentiable_error
+from ._utils.gramian import compute_gramian
+from ._utils.non_differentiable import raise_non_differentiable_error
 from .bases import _WeightedAggregator, _Weighting
 
 
diff --git a/src/torchjd/aggregation/krum.py b/src/torchjd/aggregation/krum.py
@@ -2,7 +2,7 @@
 from torch import Tensor
 from torch.nn import functional as F
 
-from ._gramian_utils import compute_gramian
+from ._utils.gramian import compute_gramian
 from .bases import _WeightedAggregator, _Weighting
 
 
diff --git a/src/torchjd/aggregation/mgda.py b/src/torchjd/aggregation/mgda.py
@@ -1,7 +1,7 @@
 import torch
 from torch import Tensor
 
-from ._gramian_utils import compute_gramian
+from ._utils.gramian import compute_gramian
 from .bases import _WeightedAggregator, _Weighting
 
 
diff --git a/src/torchjd/aggregation/nash_mtl.py b/src/torchjd/aggregation/nash_mtl.py
@@ -23,7 +23,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-from ._check_dependencies import check_dependencies_are_installed
+from ._utils.check_dependencies import check_dependencies_are_installed
 
 check_dependencies_are_installed(["cvxpy", "ecos"])
 
@@ -33,7 +33,7 @@
 from cvxpy import Expression
 from torch import Tensor
 
-from ._non_differentiable import raise_non_differentiable_error
+from ._utils.non_differentiable import raise_non_differentiable_error
 from .bases import _WeightedAggregator, _Weighting
 
 
diff --git a/src/torchjd/aggregation/pcgrad.py b/src/torchjd/aggregation/pcgrad.py
@@ -1,8 +1,8 @@
 import torch
 from torch import Tensor
 
-from ._gramian_utils import compute_gramian
-from ._non_differentiable import raise_non_differentiable_error
+from ._utils.gramian import compute_gramian
+from ._utils.non_differentiable import raise_non_differentiable_error
 from .bases import _WeightedAggregator, _Weighting
 
 
diff --git a/src/torchjd/aggregation/upgrad.py b/src/torchjd/aggregation/upgrad.py
@@ -3,10 +3,10 @@
 import torch
 from torch import Tensor
 
-from ._dual_cone_utils import project_weights
-from ._gramian_utils import compute_gramian, normalize, regularize
-from ._non_differentiable import raise_non_differentiable_error
-from ._pref_vector_utils import pref_vector_to_str_suffix, pref_vector_to_weighting
+from ._utils.dual_cone import project_weights
+from ._utils.gramian import compute_gramian, normalize, regularize
+from ._utils.non_differentiable import raise_non_differentiable_error
+from ._utils.pref_vector import pref_vector_to_str_suffix, pref_vector_to_weighting
 from .bases import _WeightedAggregator, _Weighting
 from .mean import _MeanWeighting
 
diff --git a/tests/unit/aggregation/_property_testers.py b/tests/unit/aggregation/_property_testers.py
@@ -4,7 +4,7 @@
 from torch.testing import assert_close
 
 from torchjd.aggregation import Aggregator
-from torchjd.aggregation._non_differentiable import NonDifferentiableError
+from torchjd.aggregation._utils.non_differentiable import NonDifferentiableError
 
 from ._inputs import non_strong_matrices, scaled_matrices, typical_matrices
 
diff --git a/tests/unit/aggregation/_utils/__init__.py b/tests/unit/aggregation/_utils/__init__.py
diff --git a/tests/unit/aggregation/_utils/test_dual_cone.py b/tests/unit/aggregation/_utils/test_dual_cone.py
@@ -3,7 +3,7 @@
 from pytest import mark, raises
 from torch.testing import assert_close
 
-from torchjd.aggregation._dual_cone_utils import _project_weight_vector, project_weights
+from torchjd.aggregation._utils.dual_cone import _project_weight_vector, project_weights
 
 
 @mark.parametrize("shape", [(5, 7), (9, 37), (2, 14), (32, 114), (50, 100)])
diff --git a/tests/unit/aggregation/_utils/test_pref_vector.py b/tests/unit/aggregation/_utils/test_pref_vector.py
@@ -5,7 +5,7 @@
 from torch import Tensor
 from unit._utils import ExceptionContext
 
-from torchjd.aggregation._pref_vector_utils import pref_vector_to_weighting
+from torchjd.aggregation._utils.pref_vector import pref_vector_to_weighting
 from torchjd.aggregation.mean import _MeanWeighting
 
 
