Merge branch 'main' into stationarity_property

Author: PierreQuinton

CHANGELOG.md

@@ -17,6 +17,11 @@ changes that do not affect the user.
 - Refactored the underlying optimization problem that `UPGrad` and `DualProj` have to solve to
   project onto the dual cone. This may minimally affect the output of these aggregators.
 
+### Fixed
+- Removed arbitrary exception handling in `IMTLG` and `AlignedMTL` when the computation fails. In
+  practice, this fix should only affect some matrices with extremely large values, which should
+  not usually happen.
+
 ## [0.5.0] - 2025-02-01
 
 ### Added

src/torchjd/aggregation/_gramian_utils.py

@@ -1,6 +1,5 @@
 import torch
 from torch import Tensor
-from torch.linalg import LinAlgError
 
 
 def _compute_gramian(matrix: Tensor) -> Tensor:
@@ -33,13 +32,7 @@ def _compute_normalized_gramian(matrix: Tensor, eps: float) -> Tensor:
         :math:`n` through the SVD algorithm which is efficient, therefore this is rather fast.
     """
 
-    try:
-        left_unitary_matrix, singular_values, _ = torch.linalg.svd(matrix, full_matrices=False)
-    except LinAlgError as error:  # Not sure if this can happen
-        raise ValueError(
-            f"Unexpected failure of the svd computation on matrix {matrix}. Please open an "
-            "issue on https://github.com/TorchJD/torchjd/issues and paste this error message in it."
-        ) from error
+    left_unitary_matrix, singular_values, _ = torch.linalg.svd(matrix, full_matrices=False)
     max_singular_value = torch.max(singular_values)
     if max_singular_value < eps:
         scaled_singular_values = torch.zeros_like(singular_values)

src/torchjd/aggregation/_pref_vector_utils.py

@@ -5,17 +5,6 @@
 from .constant import _ConstantWeighting
 
 
-def _check_pref_vector(pref_vector: Tensor | None) -> None:
-    """Checks the correctness of the parameter pref_vector."""
-
-    if pref_vector is not None:
-        if pref_vector.ndim != 1:
-            raise ValueError(
-                "Parameter `pref_vector` must be a vector (1D Tensor). Found `pref_vector.ndim = "
-                f"{pref_vector.ndim}`."
-            )
-
-
 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
@@ -25,6 +14,11 @@ def _pref_vector_to_weighting(pref_vector: Tensor | None, default: _Weighting) -
     if pref_vector is None:
         return default
     else:
+        if pref_vector.ndim != 1:
+            raise ValueError(
+                "Parameter `pref_vector` must be a vector (1D Tensor). Found `pref_vector.ndim = "
+                f"{pref_vector.ndim}`."
+            )
         return _ConstantWeighting(pref_vector)
 
 

src/torchjd/aggregation/aligned_mtl.py

@@ -27,13 +27,8 @@
 
 import torch
 from torch import Tensor
-from torch.linalg import LinAlgError
 
-from ._pref_vector_utils import (
-    _check_pref_vector,
-    _pref_vector_to_str_suffix,
-    _pref_vector_to_weighting,
-)
+from ._pref_vector_utils import _pref_vector_to_str_suffix, _pref_vector_to_weighting
 from .bases import _WeightedAggregator, _Weighting
 from .mean import _MeanWeighting
 
@@ -66,7 +61,6 @@ class AlignedMTL(_WeightedAggregator):
     """
 
     def __init__(self, pref_vector: Tensor | None = None):
-        _check_pref_vector(pref_vector)
         weighting = _pref_vector_to_weighting(pref_vector, default=_MeanWeighting())
         self._pref_vector = pref_vector
 
@@ -107,12 +101,7 @@ def forward(self, matrix: Tensor) -> Tensor:
     def _compute_balance_transformation(G: Tensor) -> Tensor:
         M = G.T @ G
 
-        try:
-            lambda_, V = torch.linalg.eigh(M, UPLO="U")  # More modern equivalent to torch.symeig
-        except LinAlgError:  # This can happen when the matrix has extremely large values
-            identity = torch.eye(len(M), dtype=M.dtype, device=M.device)
-            return identity
-
+        lambda_, V = torch.linalg.eigh(M, UPLO="U")  # More modern equivalent to torch.symeig
         tol = torch.max(lambda_) * len(M) * torch.finfo().eps
         rank = sum(lambda_ > tol)
 

src/torchjd/aggregation/config.py

@@ -28,13 +28,9 @@
 import torch
 from torch import Tensor
 
-from torchjd.aggregation._pref_vector_utils import (
-    _check_pref_vector,
-    _pref_vector_to_str_suffix,
-    _pref_vector_to_weighting,
-)
-from torchjd.aggregation.bases import Aggregator
-from torchjd.aggregation.sum import _SumWeighting
+from ._pref_vector_utils import _pref_vector_to_str_suffix, _pref_vector_to_weighting
+from .bases import Aggregator
+from .sum import _SumWeighting
 
 
 class ConFIG(Aggregator):
@@ -66,7 +62,6 @@ class ConFIG(Aggregator):
 
     def __init__(self, pref_vector: Tensor | None = None):
         super().__init__()
-        _check_pref_vector(pref_vector)
         self.weighting = _pref_vector_to_weighting(pref_vector, default=_SumWeighting())
         self._pref_vector = pref_vector
 

src/torchjd/aggregation/dualproj.py

@@ -4,11 +4,7 @@
 
 from ._dual_cone_utils import _project_weights
 from ._gramian_utils import _compute_regularized_normalized_gramian
-from ._pref_vector_utils import (
-    _check_pref_vector,
-    _pref_vector_to_str_suffix,
-    _pref_vector_to_weighting,
-)
+from ._pref_vector_utils import _pref_vector_to_str_suffix, _pref_vector_to_weighting
 from .bases import _WeightedAggregator, _Weighting
 from .mean import _MeanWeighting
 
@@ -51,7 +47,6 @@ def __init__(
         reg_eps: float = 0.0001,
         solver: Literal["quadprog"] = "quadprog",
     ):
-        _check_pref_vector(pref_vector)
         weighting = _pref_vector_to_weighting(pref_vector, default=_MeanWeighting())
         self._pref_vector = pref_vector
 

src/torchjd/aggregation/imtl_g.py

@@ -39,13 +39,9 @@ class _IMTLGWeighting(_Weighting):
 
     def forward(self, matrix: Tensor) -> Tensor:
         d = torch.linalg.norm(matrix, dim=1)
-
-        try:
-            v = torch.linalg.pinv(matrix @ matrix.T) @ d
-        except RuntimeError:  # This can happen when the matrix has extremely large values
-            v = torch.ones(matrix.shape[0], device=matrix.device, dtype=matrix.dtype)
-
+        v = torch.linalg.pinv(matrix @ matrix.T) @ d
         v_sum = v.sum()
+
         if v_sum.abs() < 1e-12:
             weights = torch.zeros_like(v)
         else:

src/torchjd/aggregation/upgrad.py

@@ -5,11 +5,7 @@
 
 from ._dual_cone_utils import _project_weights
 from ._gramian_utils import _compute_regularized_normalized_gramian
-from ._pref_vector_utils import (
-    _check_pref_vector,
-    _pref_vector_to_str_suffix,
-    _pref_vector_to_weighting,
-)
+from ._pref_vector_utils import _pref_vector_to_str_suffix, _pref_vector_to_weighting
 from .bases import _WeightedAggregator, _Weighting
 from .mean import _MeanWeighting
 
@@ -51,7 +47,6 @@ def __init__(
         reg_eps: float = 0.0001,
         solver: Literal["quadprog"] = "quadprog",
     ):
-        _check_pref_vector(pref_vector)
         weighting = _pref_vector_to_weighting(pref_vector, default=_MeanWeighting())
         self._pref_vector = pref_vector
 

tests/unit/aggregation/test_pref_vector_utils.py

@@ -0,0 +1,26 @@
+from contextlib import nullcontext as does_not_raise
+
+import torch
+from pytest import mark, raises
+from torch import Tensor
+from unit._utils import ExceptionContext
+
+from torchjd.aggregation._pref_vector_utils import _pref_vector_to_weighting
+from torchjd.aggregation.mean import _MeanWeighting
+
+
+@mark.parametrize(
+    ["pref_vector", "expectation"],
+    [
+        (None, does_not_raise()),
+        (torch.ones([]), raises(ValueError)),
+        (torch.ones([0]), does_not_raise()),
+        (torch.ones([1]), does_not_raise()),
+        (torch.ones([5]), does_not_raise()),
+        (torch.ones([1, 1]), raises(ValueError)),
+        (torch.ones([1, 1, 1]), raises(ValueError)),
+    ],
+)
+def test_pref_vector_to_weighting_check(pref_vector: Tensor | None, expectation: ExceptionContext):
+    with expectation:
+        _ = _pref_vector_to_weighting(pref_vector, default=_MeanWeighting())

tests/unit/autojac/_transform/test_tensor_dict.py

@@ -110,3 +110,20 @@ def _assert_class_checks_properly(
 
 def _make_tensor_dict(value_shapes: list[list[int]]) -> dict[Tensor, Tensor]:
     return {torch.zeros(key): torch.zeros(value) for key, value in zip(_key_shapes, value_shapes)}
+
+
+def test_immutability():
+    """Tests that it's impossible to modify an existing TensorDict."""
+
+    t = Gradients({})
+    with raises(TypeError):
+        t[torch.ones(1)] = torch.ones(1)
+
+    assert t == Gradients({})
+
+
+def test_empty_tensor_dict():
+    """Tests that it's impossible to instantiate a non-empty EmptyTensorDict."""
+
+    with raises(ValueError):
+        _ = EmptyTensorDict({torch.ones(1): torch.ones(1)})