Improve imports

Author: PierreQuinton

src/torchjd/aggregation/_aggregator_bases.py

@@ -2,9 +2,8 @@
 
 from torch import Tensor, nn
 
-from torchjd._utils import compute_gramian
+from torchjd._utils import Matrix, PSDMatrix, compute_gramian
 
-from .._utils.compute_gramian import Matrix, PSDMatrix
 from ._weighting_bases import Weighting
 
 

src/torchjd/aggregation/_aligned_mtl.py

@@ -28,7 +28,8 @@
 import torch
 from torch import Tensor
 
-from .._utils.compute_gramian import PSDMatrix
+from torchjd._utils.compute_gramian import PSDMatrix
+
 from ._aggregator_bases import GramianWeightedAggregator
 from ._mean import MeanWeighting
 from ._utils.pref_vector import pref_vector_to_str_suffix, pref_vector_to_weighting

src/torchjd/aggregation/_cagrad.py

@@ -1,6 +1,7 @@
 from typing import cast
 
-from .._utils.compute_gramian import PSDMatrix
+from torchjd._utils.compute_gramian import PSDMatrix
+
 from ._utils.check_dependencies import check_dependencies_are_installed
 from ._weighting_bases import Weighting
 

src/torchjd/aggregation/_constant.py

@@ -1,6 +1,7 @@
 from torch import Tensor
 
-from .._utils.compute_gramian import Matrix
+from torchjd._utils.compute_gramian import Matrix
+
 from ._aggregator_bases import WeightedAggregator
 from ._utils.str import vector_to_str
 from ._weighting_bases import Weighting

src/torchjd/aggregation/_dualproj.py

@@ -2,7 +2,8 @@
 
 from torch import Tensor
 
-from .._utils.compute_gramian import PSDMatrix
+from torchjd._utils.compute_gramian import PSDMatrix
+
 from ._aggregator_bases import GramianWeightedAggregator
 from ._mean import MeanWeighting
 from ._utils.dual_cone import project_weights

src/torchjd/aggregation/_imtl_g.py

@@ -1,7 +1,8 @@
 import torch
 from torch import Tensor
 
-from .._utils.compute_gramian import PSDMatrix
+from torchjd._utils.compute_gramian import PSDMatrix
+
 from ._aggregator_bases import GramianWeightedAggregator
 from ._utils.non_differentiable import raise_non_differentiable_error
 from ._weighting_bases import Weighting

src/torchjd/aggregation/_krum.py

@@ -2,7 +2,8 @@
 from torch import Tensor
 from torch.nn import functional as F
 
-from .._utils.compute_gramian import PSDMatrix
+from torchjd._utils.compute_gramian import PSDMatrix
+
 from ._aggregator_bases import GramianWeightedAggregator
 from ._weighting_bases import Weighting
 

src/torchjd/aggregation/_mean.py

@@ -1,7 +1,8 @@
 import torch
 from torch import Tensor
 
-from .._utils.compute_gramian import Matrix
+from torchjd._utils.compute_gramian import Matrix
+
 from ._aggregator_bases import WeightedAggregator
 from ._weighting_bases import Weighting
 

src/torchjd/aggregation/_mgda.py

@@ -1,7 +1,8 @@
 import torch
 from torch import Tensor
 
-from .._utils.compute_gramian import PSDMatrix
+from torchjd._utils.compute_gramian import PSDMatrix
+
 from ._aggregator_bases import GramianWeightedAggregator
 from ._weighting_bases import Weighting
 

src/torchjd/aggregation/_nash_mtl.py

@@ -25,7 +25,8 @@
 
 # mypy: ignore-errors
 
-from .._utils.compute_gramian import Matrix
+from torchjd._utils.compute_gramian import Matrix
+
 from ._utils.check_dependencies import check_dependencies_are_installed
 from ._weighting_bases import Weighting