Replace check_and_get_keys by check_keys

Author: PierreQuinton

src/torchjd/autojac/_transform/_differentiate.py

@@ -3,7 +3,7 @@
 
 from torch import Tensor
 
-from .base import _A, Transform
+from .base import _A, RequirementError, Transform
 from .ordered_set import OrderedSet
 
 
@@ -38,6 +38,11 @@ def _differentiate(self, tensor_outputs: Sequence[Tensor]) -> tuple[Tensor, ...]
         tensor_outputs should be.
         """
 
-    def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
-        # outputs in the forward direction become inputs in the backward direction, and vice-versa
-        return set(self.outputs), set(self.inputs)
+    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+        outputs = set(self.outputs)
+        if not outputs.issubset(input_keys):
+            raise RequirementError(
+                f"The input_keys needs to be a super set of the outputs. Found {input_keys} and "
+                f"{outputs}"
+            )
+        return set(self.inputs)

src/torchjd/autojac/_transform/accumulate.py

@@ -28,8 +28,13 @@ def __call__(self, gradients: Gradients) -> EmptyTensorDict:
 
         return EmptyTensorDict()
 
-    def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
-        return self._required_keys, set()
+    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()
 
 
 def _check_expects_grad(tensor: Tensor) -> None:

src/torchjd/autojac/_transform/aggregate.py

@@ -6,7 +6,7 @@
 
 from torchjd.aggregation import Aggregator
 
-from .base import Transform
+from .base import RequirementError, Transform
 from .ordered_set import OrderedSet
 from .tensor_dict import EmptyTensorDict, Gradients, GradientVectors, JacobianMatrices, Jacobians
 
@@ -16,18 +16,18 @@
 
 class Aggregate(Transform[Jacobians, Gradients]):
     def __init__(self, aggregator: Aggregator, key_order: Iterable[Tensor]):
-        matrixify = _Matrixify(key_order)
+        matrixify = _Matrixify()
         aggregate_matrices = _AggregateMatrices(aggregator, key_order)
-        reshape = _Reshape(key_order)
+        reshape = _Reshape()
 
         self._aggregator_str = str(aggregator)
         self.transform = reshape << aggregate_matrices << matrixify
 
     def __call__(self, input: Jacobians) -> Gradients:
         return self.transform(input)
 
-    def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
-        return self.transform.check_and_get_keys()
+    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+        return self.transform.check_keys(input_keys)
 
 
 class _AggregateMatrices(Transform[JacobianMatrices, GradientVectors]):
@@ -48,9 +48,12 @@ def __call__(self, jacobian_matrices: JacobianMatrices) -> GradientVectors:
         ordered_matrices = self._select_ordered_subdict(jacobian_matrices, self.key_order)
         return self._aggregate_group(ordered_matrices, self.aggregator)
 
-    def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
-        keys = set(self.key_order)
-        return keys, keys
+    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+        if not set(self.key_order) == input_keys:
+            raise RequirementError(
+                f"The input_keys must match the key_order. Found {input_keys} and {self.key_order}"
+            )
+        return input_keys
 
     @staticmethod
     def _select_ordered_subdict(
@@ -108,29 +111,23 @@ def _disunite(
 
 
 class _Matrixify(Transform[Jacobians, JacobianMatrices]):
-    def __init__(self, required_keys: Iterable[Tensor]):
-        self._required_keys = set(required_keys)
-
     def __call__(self, jacobians: Jacobians) -> JacobianMatrices:
         jacobian_matrices = {
             key: jacobian.view(jacobian.shape[0], -1) for key, jacobian in jacobians.items()
         }
         return JacobianMatrices(jacobian_matrices)
 
-    def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
-        return self._required_keys, self._required_keys
+    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+        return input_keys
 
 
 class _Reshape(Transform[GradientVectors, Gradients]):
-    def __init__(self, required_keys: Iterable[Tensor]):
-        self._required_keys = set(required_keys)
-
     def __call__(self, gradient_vectors: GradientVectors) -> Gradients:
         gradients = {
             key: gradient_vector.view(key.shape)
             for key, gradient_vector in gradient_vectors.items()
         }
         return Gradients(gradients)
 
-    def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
-        return self._required_keys, self._required_keys
+    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+        return input_keys

src/torchjd/autojac/_transform/base.py

@@ -48,17 +48,13 @@ def __call__(self, input: _B) -> _C:
         """Applies the transform to the input."""
 
     @abstractmethod
-    def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
+    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
         """
-        Returns a pair containing (in order) the required keys and the output keys of the Transform
-        and recursively checks that the transform is valid.
+        Checks the keys of the Transform for the provided input_keys and returns the corresponding
+        output keys for recursion.
 
-        The required keys are the set of keys that the transform requires to be present in its input
-        TensorDicts. The output keys are the set of keys that will be present in the output
-        TensorDicts of the transform.
-
-        Since the computation of the required and output keys and the verification that the
-        transform is valid are sometimes intertwined operations, we do them in a single method.
+        The output keys are the set of keys that will be present in the output TensorDict of the
+        transform given that the provided TensorDict has the provided input_keys.
         """
 
     __lshift__ = compose
@@ -77,15 +73,9 @@ def __call__(self, input: _A) -> _C:
         intermediate = self.inner(input)
         return self.outer(intermediate)
 
-    def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
-        outer_required_keys, outer_output_keys = self.outer.check_and_get_keys()
-        inner_required_keys, inner_output_keys = self.inner.check_and_get_keys()
-        if outer_required_keys != inner_output_keys:
-            raise RequirementError(
-                "The `output_keys` of `inner` must match with the `required_keys` of "
-                f"outer. Found {outer_required_keys} and {inner_output_keys}"
-            )
-        return inner_required_keys, outer_output_keys
+    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+        intermediate_keys = self.inner.check_keys(input_keys)
+        return self.outer.check_keys(intermediate_keys)
 
 
 class Conjunction(Transform[_A, _B]):
@@ -106,18 +96,11 @@ def __call__(self, tensor_dict: _A) -> _B:
         output = _union([transform(tensor_dict) for transform in self.transforms])
         return output
 
-    def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
-        keys_pairs = [transform.check_and_get_keys() for transform in self.transforms]
-
-        required_keys = set(key for required_keys, _ in keys_pairs for key in required_keys)
-        for transform_required_keys, _ in keys_pairs:
-            if transform_required_keys != required_keys:
-                raise RequirementError("All transforms should require the same set of keys.")
-
-        output_keys_with_duplicates = [key for _, output_keys in keys_pairs for key in output_keys]
-        output_keys = set(output_keys_with_duplicates)
+    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+        output_keys_list = [key for t in self.transforms for key in t.check_keys(input_keys)]
+        output_keys = set(output_keys_list)
 
-        if len(output_keys) != len(output_keys_with_duplicates):
+        if len(output_keys) != len(output_keys_list):
             raise RequirementError("The sets of output keys of transforms should be disjoint.")
 
-        return required_keys, output_keys
+        return output_keys

src/torchjd/autojac/_transform/diagonalize.py

@@ -3,7 +3,7 @@
 import torch
 from torch import Tensor
 
-from .base import Transform
+from .base import RequirementError, Transform
 from .ordered_set import OrderedSet
 from .tensor_dict import Gradients, Jacobians
 
@@ -27,6 +27,11 @@ def __call__(self, tensors: Gradients) -> Jacobians:
         }
         return Jacobians(diagonalized_tensors)
 
-    def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
-        keys = set(self.considered)
-        return keys, keys
+    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+        considered = set(self.considered)
+        if not considered.issubset(input_keys):
+            raise RequirementError(
+                f"The input_keys needs to be a super set of the considered keys. Found {input_keys} "
+                f"and {considered}"
+            )
+        return considered

src/torchjd/autojac/_transform/init.py

@@ -3,7 +3,7 @@
 import torch
 from torch import Tensor
 
-from .base import Transform
+from .base import RequirementError, Transform
 from .tensor_dict import EmptyTensorDict, Gradients
 
 
@@ -21,5 +21,7 @@ def __call__(self, input: EmptyTensorDict) -> Gradients:
 
         return Gradients({value: torch.ones_like(value) for value in self.values})
 
-    def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
-        return set(), self.values
+    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+        if input_keys == set():
+            raise RequirementError(f"Init expects an empty set of input_keys. Found {input_keys}")
+        return self.values

src/torchjd/autojac/_transform/select.py

@@ -7,19 +7,17 @@
 
 
 class Select(Transform[_A, _A]):
-    def __init__(self, keys: Iterable[Tensor], required_keys: Iterable[Tensor]):
+    def __init__(self, keys: Iterable[Tensor]):
         self.keys = set(keys)
-        self._required_keys = set(required_keys)
 
     def __call__(self, tensor_dict: _A) -> _A:
         output = {key: tensor_dict[key] for key in self.keys}
         return type(tensor_dict)(output)
 
-    def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
-        required_keys = self._required_keys
-        if not self.keys.issubset(required_keys):
+    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+        if not self.keys.issubset(input_keys):
             raise RequirementError(
-                "Parameter `keys` should be a subset of parameter `required_keys`"
+                f"The input_keys needs to be a super set of the keys to select. Found {input_keys} "
+                f"and {self.keys}"
             )
-
-        return required_keys, self.keys
+        return self.keys

src/torchjd/autojac/_transform/stack.py

@@ -4,7 +4,7 @@
 from torch import Tensor
 
 from ._utils import _A, _materialize, dicts_union
-from .base import RequirementError, Transform
+from .base import Transform
 from .tensor_dict import Gradients, Jacobians
 
 
@@ -17,17 +17,8 @@ def __call__(self, input: _A) -> Jacobians:
         result = _stack(results)
         return result
 
-    def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
-        keys_pairs = [transform.check_and_get_keys() for transform in self.transforms]
-
-        required_keys = set(key for required_keys, _ in keys_pairs for key in required_keys)
-        output_keys = set(key for _, output_keys in keys_pairs for key in output_keys)
-
-        for transform_required_keys, _ in keys_pairs:
-            if transform_required_keys != required_keys:
-                raise RequirementError("All transforms should require the same set of keys.")
-
-        return required_keys, output_keys
+    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+        return {key for transform in self.transforms for key in transform.check_keys(input_keys)}
 
 
 def _stack(gradient_dicts: list[Gradients]) -> Jacobians:

src/torchjd/autojac/mtl_backward.py

@@ -179,10 +179,10 @@ 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, to_differentiate)
+    accumulate = Accumulate(task_params) << Select(task_params)
 
     # Transform that backpropagates the gradients of the losses w.r.t. the features.
-    backpropagate = Select(features, to_differentiate)
+    backpropagate = Select(features)
 
     # Transform that accumulates the gradient of the losses w.r.t. the task-specific parameters into
     # their .grad fields and backpropagates the gradient of the losses w.r.t. to the features.