Make Accumulate not check required_keys

Author: ValerianRey

src/torchjd/autojac/_transform/accumulate.py

@@ -1,15 +1,10 @@
-from typing import Iterable
-
 from torch import Tensor
 
-from .base import RequirementError, Transform
+from .base import Transform
 from .tensor_dict import EmptyTensorDict, Gradients
 
 
 class Accumulate(Transform[Gradients, EmptyTensorDict]):
-    def __init__(self, required_keys: Iterable[Tensor]):
-        self._required_keys = set(required_keys)
-
     def __call__(self, gradients: Gradients) -> EmptyTensorDict:
         """
         Accumulates gradients with respect to keys in their ``.grad`` field.
@@ -29,11 +24,6 @@ def __call__(self, gradients: Gradients) -> EmptyTensorDict:
         return EmptyTensorDict()
 
     def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
-        if not self._required_keys.issubset(input_keys):
-            raise RequirementError(
-                f"The input_keys needs to be a super set of the required_keys. Found {input_keys} "
-                f"and {self._required_keys}"
-            )
         return set()
 
 

src/torchjd/autojac/backward.py

@@ -112,6 +112,6 @@ def _create_transform(
     aggregate = Aggregate(aggregator, inputs)
 
     # Transform that accumulates the result in the .grad field of the inputs.
-    accumulate = Accumulate(inputs)
+    accumulate = Accumulate()
 
     return accumulate << aggregate << jac << diag << init

src/torchjd/autojac/mtl_backward.py

@@ -155,7 +155,7 @@ def _create_transform(
     aggregate = Aggregate(aggregator, shared_params)
 
     # Transform that accumulates the result in the .grad field of the shared parameters.
-    accumulate = Accumulate(shared_params)
+    accumulate = Accumulate()
 
     return accumulate << aggregate << jac << stack
 
@@ -179,7 +179,7 @@ def _create_task_transform(
 
     # Transform that accumulates the gradients w.r.t. the task-specific parameters into their
     # .grad fields.
-    accumulate = Accumulate(task_params) << Select(task_params)
+    accumulate = Accumulate() << Select(task_params)
 
     # Transform that backpropagates the gradients of the losses w.r.t. the features.
     backpropagate = Select(features)