fix(aggregation): Make non-differentiability explicit (#334)

* Add _non_differentiable.py with NonDifferentiableError and raise_non_differentiable_error
* Register raise_non_differentiable_error as a full backward pre-hook of CAGrad, ConFIG, DualProj, GradDrop, IMTLG, NashMTL, PCGrad and UPGrad
* Add NonDifferentiableProperty tester
* Give NonDifferentiableProperty to CAGrad, ConFIG, DualProj, GradDrop, IMTLG, NashMTL, PCGrad and UPGrad
* Add changelog entry

Author: ValerianRey

CHANGELOG.md

@@ -15,6 +15,13 @@ changes that do not affect the user.
   torchjd[nash_mtl]` or `pip install torchjd[full]` to install TorchJD alongside those dependencies.
   This should make TorchJD more lightweight.
 
+### Fixed
+
+- Made some aggregators (`CAGrad`, `ConFIG`, `DualProj`, `GradDrop`, `IMTLG`, `NashMTL`, `PCGrad`
+  and `UPGrad`) raise a `NonDifferentiableError` whenever one tries to differentiate through them.
+  Before this change, trying to differentiate through them leaded to wrong gradients or unclear
+  errors.
+
 ## [0.6.0] - 2025-04-19
 
 ### Added

src/torchjd/aggregation/_non_differentiable.py

@@ -0,0 +1,10 @@
+from torch import Tensor, nn
+
+
+class NonDifferentiableError(RuntimeError):
+    def __init__(self, module: nn.Module):
+        super().__init__(f"Trying to differentiate through {module}, which is not differentiable.")
+
+
+def raise_non_differentiable_error(module: nn.Module, _: tuple[Tensor, ...]) -> None:
+    raise NonDifferentiableError(module)

src/torchjd/aggregation/cagrad.py

@@ -8,6 +8,7 @@
 from torch import Tensor
 
 from ._gramian_utils import compute_gramian, normalize
+from ._non_differentiable import raise_non_differentiable_error
 from .bases import _WeightedAggregator, _Weighting
 
 
@@ -44,6 +45,9 @@ class CAGrad(_WeightedAggregator):
     def __init__(self, c: float, norm_eps: float = 0.0001):
         super().__init__(weighting=_CAGradWeighting(c=c, norm_eps=norm_eps))
 
+        # This prevents considering the computed weights as constant w.r.t. the matrix.
+        self.register_full_backward_pre_hook(raise_non_differentiable_error)
+
     def __repr__(self) -> str:
         return (
             f"{self.__class__.__name__}(c={self.weighting.c}, norm_eps={self.weighting.norm_eps})"

src/torchjd/aggregation/config.py

@@ -28,6 +28,7 @@
 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 .bases import Aggregator
 from .sum import _SumWeighting
@@ -65,6 +66,9 @@ def __init__(self, pref_vector: Tensor | None = None):
         self.weighting = pref_vector_to_weighting(pref_vector, default=_SumWeighting())
         self._pref_vector = pref_vector
 
+        # This prevents computing gradients that can be very wrong.
+        self.register_full_backward_pre_hook(raise_non_differentiable_error)
+
     def forward(self, matrix: Tensor) -> Tensor:
         weights = self.weighting(matrix)
         units = torch.nan_to_num((matrix / (matrix.norm(dim=1)).unsqueeze(1)), 0.0)

src/torchjd/aggregation/dualproj.py

@@ -4,6 +4,7 @@
 
 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 .bases import _WeightedAggregator, _Weighting
 from .mean import _MeanWeighting
@@ -56,6 +57,9 @@ def __init__(
             )
         )
 
+        # This prevents considering the computed weights as constant w.r.t. the matrix.
+        self.register_full_backward_pre_hook(raise_non_differentiable_error)
+
     def __repr__(self) -> str:
         return (
             f"{self.__class__.__name__}(pref_vector={repr(self._pref_vector)}, norm_eps="

src/torchjd/aggregation/graddrop.py

@@ -3,6 +3,7 @@
 import torch
 from torch import Tensor
 
+from ._non_differentiable import raise_non_differentiable_error
 from .bases import Aggregator
 
 
@@ -47,6 +48,9 @@ def __init__(self, f: Callable = _identity, leak: Tensor | None = None):
         self.f = f
         self.leak = leak
 
+        # This prevents computing gradients that can be very wrong.
+        self.register_full_backward_pre_hook(raise_non_differentiable_error)
+
     def forward(self, matrix: Tensor) -> Tensor:
         self._check_is_matrix(matrix)
         self._check_matrix_has_enough_rows(matrix)

src/torchjd/aggregation/imtl_g.py

@@ -2,6 +2,7 @@
 from torch import Tensor
 
 from ._gramian_utils import compute_gramian
+from ._non_differentiable import raise_non_differentiable_error
 from .bases import _WeightedAggregator, _Weighting
 
 
@@ -30,6 +31,9 @@ class IMTLG(_WeightedAggregator):
     def __init__(self):
         super().__init__(weighting=_IMTLGWeighting())
 
+        # This prevents computing gradients that can be very wrong.
+        self.register_full_backward_pre_hook(raise_non_differentiable_error)
+
 
 class _IMTLGWeighting(_Weighting):
     """

src/torchjd/aggregation/nash_mtl.py

@@ -33,6 +33,7 @@
 from cvxpy import Expression
 from torch import Tensor
 
+from ._non_differentiable import raise_non_differentiable_error
 from .bases import _WeightedAggregator, _Weighting
 
 
@@ -95,6 +96,9 @@ def __init__(
             )
         )
 
+        # This prevents considering the computed weights as constant w.r.t. the matrix.
+        self.register_full_backward_pre_hook(raise_non_differentiable_error)
+
     def reset(self) -> None:
         """Resets the internal state of the algorithm."""
         self.weighting.reset()

src/torchjd/aggregation/pcgrad.py

@@ -2,6 +2,7 @@
 from torch import Tensor
 
 from ._gramian_utils import compute_gramian
+from ._non_differentiable import raise_non_differentiable_error
 from .bases import _WeightedAggregator, _Weighting
 
 
@@ -28,6 +29,9 @@ class PCGrad(_WeightedAggregator):
     def __init__(self):
         super().__init__(weighting=_PCGradWeighting())
 
+        # This prevents running into a RuntimeError due to modifying stored tensors in place.
+        self.register_full_backward_pre_hook(raise_non_differentiable_error)
+
 
 class _PCGradWeighting(_Weighting):
     """

src/torchjd/aggregation/upgrad.py

@@ -5,6 +5,7 @@
 
 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 .bases import _WeightedAggregator, _Weighting
 from .mean import _MeanWeighting
@@ -56,6 +57,9 @@ def __init__(
             )
         )
 
+        # This prevents considering the computed weights as constant w.r.t. the matrix.
+        self.register_full_backward_pre_hook(raise_non_differentiable_error)
+
     def __repr__(self) -> str:
         return (
             f"{self.__class__.__name__}(pref_vector={repr(self._pref_vector)}, norm_eps="