refactor(aggregation): Remove _check_is_finite (#381)

* Remove _check_is_finite
* Add changelog entry

Author: ValerianRey

CHANGELOG.md

@@ -21,6 +21,10 @@ changes that do not affect the user.
   always be imported directly from the `torchjd` package (e.g.
   `from torchjd.autojac.mtl_backward import mtl_backward` must be changed to
   `from torchjd import mtl_backward`).
+- Removed the check that the input Jacobian matrix provided to an aggregator does not contain `nan`,
+  `inf` or `-inf` values. This check was costly in memory and in time for large matrices so this
+  should improve performance. However, if the optimization diverges for some reason (for instance
+  due to a too large learning rate), the resulting exceptions may come from other sources.
 
 ### Fixed
 

src/torchjd/aggregation/_aggregator_bases.py

@@ -20,14 +20,6 @@ def _check_is_matrix(matrix: Tensor) -> None:
                 f"{matrix.shape}`."
             )
 
-    @staticmethod
-    def _check_is_finite(matrix: Tensor) -> None:
-        if not matrix.isfinite().all():
-            raise ValueError(
-                "Parameter `matrix` should be a tensor of finite elements (no nan, inf or -inf "
-                f"values). Found `matrix = {matrix}`."
-            )
-
     @abstractmethod
     def forward(self, matrix: Tensor) -> Tensor:
         """Computes the aggregation from the input matrix."""
@@ -69,8 +61,6 @@ def combine(matrix: Tensor, weights: Tensor) -> Tensor:
 
     def forward(self, matrix: Tensor) -> Tensor:
         self._check_is_matrix(matrix)
-        self._check_is_finite(matrix)
-
         weights = self.weighting(matrix)
         vector = self.combine(matrix, weights)
         return vector

src/torchjd/aggregation/_graddrop.py

@@ -55,7 +55,6 @@ def __init__(self, f: Callable = _identity, leak: Tensor | None = None):
     def forward(self, matrix: Tensor) -> Tensor:
         self._check_is_matrix(matrix)
         self._check_matrix_has_enough_rows(matrix)
-        self._check_is_finite(matrix)
 
         if matrix.shape[0] == 0 or matrix.shape[1] == 0:
             return torch.zeros(matrix.shape[1], dtype=matrix.dtype, device=matrix.device)

src/torchjd/aggregation/_trimmed_mean.py

@@ -47,7 +47,6 @@ def __init__(self, trim_number: int):
     def forward(self, matrix: Tensor) -> Tensor:
         self._check_is_matrix(matrix)
         self._check_matrix_has_enough_rows(matrix)
-        self._check_is_finite(matrix)
 
         n_rows = matrix.shape[0]
         n_remaining = n_rows - 2 * self.trim_number

tests/unit/aggregation/test_aggregator_bases.py

@@ -21,19 +21,3 @@
 def test_check_is_matrix(shape: Sequence[int], expectation: ExceptionContext):
     with expectation:
         Aggregator._check_is_matrix(torch.randn(shape))
-
-
-@mark.parametrize(
-    ["value", "expectation"],
-    [
-        (0.0, does_not_raise()),
-        (torch.nan, raises(ValueError)),
-        (torch.inf, raises(ValueError)),
-        (-torch.inf, raises(ValueError)),
-    ],
-)
-def test_check_is_finite(value: float, expectation: ExceptionContext):
-    matrix = torch.ones([5, 5])
-    matrix[1, 2] = value
-    with expectation:
-        Aggregator._check_is_finite(matrix)