refactor(autojac): Remove _disunite check (#377)

Author: ValerianRey

src/torchjd/_autojac/_transform/_aggregate.py

@@ -115,15 +115,6 @@ def _unite(jacobian_matrices: OrderedDict[Tensor, Tensor]) -> Tensor:
     def _disunite(
         united_gradient_vector: Tensor, jacobian_matrices: OrderedDict[Tensor, Tensor]
     ) -> GradientVectors:
-        expected_length = sum([matrix.shape[1] for matrix in jacobian_matrices.values()])
-        if len(united_gradient_vector) != expected_length:
-            raise ValueError(
-                "Parameter `united_gradient_vector` should be a vector with length equal to the sum"
-                "of the numbers of columns in the jacobian matrices. Found"
-                f"`len(united_gradient_vector) = {len(united_gradient_vector)}` and the sum of the "
-                f"numbers of columns in the jacobian matrices is {expected_length}."
-            )
-
         gradient_vectors = {}
         start = 0
         for key, jacobian_matrix in jacobian_matrices.items():

tests/unit/autojac/_transform/test_aggregate.py

@@ -1,9 +1,7 @@
 import math
-from collections import OrderedDict
 
 import torch
 from pytest import mark, raises
-from torch import Tensor
 from unit.conftest import DEVICE
 
 from torchjd._autojac._transform import (
@@ -72,37 +70,6 @@ def test_aggregate_matrices_empty_dict():
     assert len(gradient_vectors) == 0
 
 
-@mark.parametrize(
-    ["united_gradient_vector", "jacobian_matrices"],
-    [
-        (
-            torch.ones(10),
-            {  # Total number of parameters according to the united gradient vector: 10
-                torch.ones(5): torch.ones(2, 5),
-                torch.ones(4): torch.ones(2, 4),
-            },
-        ),  # Total number of parameters according to the jacobian matrices: 9
-        (
-            torch.ones(10),
-            {  # Total number of parameters according to the united gradient vector: 10
-                torch.ones(5): torch.ones(2, 5),
-                torch.ones(3): torch.ones(2, 3),
-                torch.ones(3): torch.ones(2, 3),
-            },
-        ),  # Total number of parameters according to the jacobian matrices: 11
-    ],
-)
-def test_disunite_wrong_vector_length(
-    united_gradient_vector: Tensor, jacobian_matrices: dict[Tensor, Tensor]
-):
-    """
-    Tests that the _disunite method raises a ValueError when used on vectors of the wrong length.
-    """
-
-    with raises(ValueError):
-        _AggregateMatrices._disunite(united_gradient_vector, OrderedDict(jacobian_matrices))
-
-
 def test_matrixify():
     """Tests that the Matrixify transform correctly creates matrices from the jacobians."""