Isolate key checks in transforms

src/torchjd/autojac/_transform/_differentiate.py

@@ -38,12 +38,6 @@ def _differentiate(self, tensor_outputs: Sequence[Tensor]) -> tuple[Tensor, ...]
         tensor_outputs should be.
         """
 
-    @property
-    def required_keys(self) -> set[Tensor]:
+    def check_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)
-
-    @property
-    def output_keys(self) -> set[Tensor]:
-        # outputs in the forward direction become inputs in the backward direction, and vice-versa
-        return set(self.inputs)
+        return set(self.outputs), set(self.inputs)

src/torchjd/autojac/_transform/accumulate.py

@@ -28,13 +28,8 @@ def _compute(self, gradients: Gradients) -> EmptyTensorDict:
 
         return EmptyTensorDict()
 
-    @property
-    def required_keys(self) -> set[Tensor]:
-        return self._required_keys
-
-    @property
-    def output_keys(self) -> set[Tensor]:
-        return set()
+    def check_keys(self) -> tuple[set[Tensor], set[Tensor]]:
+        return self._required_keys, set()
 
 
 def _check_expects_grad(tensor: Tensor) -> None:

src/torchjd/autojac/_transform/aggregate.py

@@ -26,13 +26,8 @@ def __init__(self, aggregator: Aggregator, key_order: Iterable[Tensor]):
     def _compute(self, input: Jacobians) -> Gradients:
         return self.transform(input)
 
-    @property
-    def required_keys(self) -> set[Tensor]:
-        return self.transform.required_keys
-
-    @property
-    def output_keys(self) -> set[Tensor]:
-        return self.transform.output_keys
+    def check_keys(self) -> tuple[set[Tensor], set[Tensor]]:
+        return self.transform.check_keys()
 
 
 class _AggregateMatrices(Transform[JacobianMatrices, GradientVectors]):
@@ -53,13 +48,9 @@ def _compute(self, jacobian_matrices: JacobianMatrices) -> GradientVectors:
         ordered_matrices = self._select_ordered_subdict(jacobian_matrices, self.key_order)
         return self._aggregate_group(ordered_matrices, self.aggregator)
 
-    @property
-    def required_keys(self) -> set[Tensor]:
-        return set(self.key_order)
-
-    @property
-    def output_keys(self) -> set[Tensor]:
-        return set(self.key_order)
+    def check_keys(self) -> tuple[set[Tensor], set[Tensor]]:
+        keys = set(self.key_order)
+        return keys, keys
 
     @staticmethod
     def _select_ordered_subdict(
@@ -126,13 +117,8 @@ def _compute(self, jacobians: Jacobians) -> JacobianMatrices:
         }
         return JacobianMatrices(jacobian_matrices)
 
-    @property
-    def required_keys(self) -> set[Tensor]:
-        return self._required_keys
-
-    @property
-    def output_keys(self) -> set[Tensor]:
-        return self._required_keys
+    def check_keys(self) -> tuple[set[Tensor], set[Tensor]]:
+        return self._required_keys, self._required_keys
 
 
 class _Reshape(Transform[GradientVectors, Gradients]):
@@ -146,10 +132,5 @@ def _compute(self, gradient_vectors: GradientVectors) -> Gradients:
         }
         return Gradients(gradients)
 
-    @property
-    def required_keys(self) -> set[Tensor]:
-        return self._required_keys
-
-    @property
-    def output_keys(self) -> set[Tensor]:
-        return self._required_keys
+    def check_keys(self) -> tuple[set[Tensor], set[Tensor]]:
+        return self._required_keys, self._required_keys

src/torchjd/autojac/_transform/base.py

@@ -44,32 +44,21 @@ def _compute(self, input: _B) -> _C:
         """Applies the transform to the input."""
 
     def __call__(self, input: _B) -> _C:
-        input.check_keys_are(self.required_keys)
         return self._compute(input)
 
-    @property
     @abstractmethod
-    def required_keys(self) -> set[Tensor]:
+    def check_keys(self) -> tuple[set[Tensor], set[Tensor]]:
         """
-        Returns the set of keys that the transform requires to be present in its input TensorDicts.
+        Returns a pair containing (in order) the required keys and the output keys of the Transform.
+        Checks that the transform is valid.
         """
 
-    @property
-    @abstractmethod
-    def output_keys(self) -> set[Tensor]:
-        """Returns the set of keys that will be present in the output of the transform."""
-
     __lshift__ = compose
     __or__ = conjunct
 
 
 class Composition(Transform[_A, _C]):
     def __init__(self, outer: Transform[_B, _C], inner: Transform[_A, _B]):
-        if outer.required_keys != inner.output_keys:
-            raise ValueError(
-                "The `output_keys` of `inner` must match with the `required_keys` of "
-                f"outer. Found {outer.required_keys} and {inner.output_keys}"
-            )
         self.outer = outer
         self.inner = inner
 
@@ -80,32 +69,21 @@ def _compute(self, input: _A) -> _C:
         intermediate = self.inner(input)
         return self.outer(intermediate)
 
-    @property
-    def required_keys(self) -> set[Tensor]:
-        return self.inner.required_keys
-
-    @property
-    def output_keys(self) -> set[Tensor]:
-        return self.outer.output_keys
+    def check_keys(self) -> tuple[set[Tensor], set[Tensor]]:
+        outer_required_keys, outer_output_keys = self.outer.check_keys()
+        inner_required_keys, inner_output_keys = self.inner.check_keys()
+        if outer_required_keys != inner_output_keys:
+            raise ValueError(
+                "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
 
 
 class Conjunction(Transform[_A, _B]):
     def __init__(self, transforms: Sequence[Transform[_A, _B]]):
         self.transforms = transforms
 
-        self._required_keys = set(
-            key for transform in transforms for key in transform.required_keys
-        )
-        for transform in transforms:
-            if transform.required_keys != self.required_keys:
-                raise ValueError("All transforms should require the same set of keys.")
-
-        output_keys_with_duplicates = [key for t in transforms for key in t.output_keys]
-        self._output_keys = set(output_keys_with_duplicates)
-
-        if len(self._output_keys) != len(output_keys_with_duplicates):
-            raise ValueError("The sets of output keys of transforms should be disjoint.")
-
     def __str__(self) -> str:
         strings = []
         for t in self.transforms:
@@ -120,10 +98,18 @@ def _compute(self, tensor_dict: _A) -> _B:
         output = _union([transform(tensor_dict) for transform in self.transforms])
         return output
 
-    @property
-    def required_keys(self) -> set[Tensor]:
-        return self._required_keys
+    def check_keys(self) -> tuple[set[Tensor], set[Tensor]]:
+        keys_pairs = [transform.check_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 ValueError("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)
+
+        if len(output_keys) != len(output_keys_with_duplicates):
+            raise ValueError("The sets of output keys of transforms should be disjoint.")
 
-    @property
-    def output_keys(self) -> set[Tensor]:
-        return self._output_keys
+        return required_keys, output_keys

src/torchjd/autojac/_transform/diagonalize.py

@@ -27,10 +27,6 @@ def _compute(self, tensors: Gradients) -> Jacobians:
         }
         return Jacobians(diagonalized_tensors)
 
-    @property
-    def required_keys(self) -> set[Tensor]:
-        return set(self.considered)
-
-    @property
-    def output_keys(self) -> set[Tensor]:
-        return set(self.considered)
+    def check_keys(self) -> tuple[set[Tensor], set[Tensor]]:
+        keys = set(self.considered)
+        return keys, keys

src/torchjd/autojac/_transform/init.py

@@ -21,10 +21,5 @@ def _compute(self, input: EmptyTensorDict) -> Gradients:
 
         return Gradients({value: torch.ones_like(value) for value in self.values})
 
-    @property
-    def required_keys(self) -> set[Tensor]:
-        return set()
-
-    @property
-    def output_keys(self) -> set[Tensor]:
-        return self.values
+    def check_keys(self) -> tuple[set[Tensor], set[Tensor]]:
+        return set(), self.values

src/torchjd/autojac/_transform/select.py

@@ -18,10 +18,5 @@ def _compute(self, tensor_dict: _A) -> _A:
         output = {key: tensor_dict[key] for key in self.keys}
         return type(tensor_dict)(output)
 
-    @property
-    def required_keys(self) -> set[Tensor]:
-        return self._required_keys
-
-    @property
-    def output_keys(self) -> set[Tensor]:
-        return self.keys
+    def check_keys(self) -> tuple[set[Tensor], set[Tensor]]:
+        return self._required_keys, self.keys

src/torchjd/autojac/_transform/stack.py

@@ -12,25 +12,22 @@ class Stack(Transform[_A, Jacobians]):
     def __init__(self, transforms: Sequence[Transform[_A, Gradients]]):
         self.transforms = transforms
 
-        self._required_keys = {key for transform in transforms for key in transform.required_keys}
-        self._output_keys = {key for transform in transforms for key in transform.output_keys}
-
-        for transform in transforms:
-            if transform.required_keys != self.required_keys:
-                raise ValueError("All transforms should require the same set of keys.")
-
     def _compute(self, input: _A) -> Jacobians:
         results = [transform(input) for transform in self.transforms]
         result = _stack(results)
         return result
 
-    @property
-    def required_keys(self) -> set[Tensor]:
-        return self._required_keys
+    def check_keys(self) -> tuple[set[Tensor], set[Tensor]]:
+        keys_pairs = [transform.check_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 ValueError("All transforms should require the same set of keys.")
 
-    @property
-    def output_keys(self) -> set[Tensor]:
-        return self._output_keys
+        return required_keys, output_keys
 
 
 def _stack(gradient_dicts: list[Gradients]) -> Jacobians:

tests/unit/autojac/_transform/test_base.py

@@ -27,13 +27,8 @@ def _compute(self, input: _B) -> _C:
         output_dict = {key: torch.empty(0) for key in self._output_keys}
         return typing.cast(_C, output_dict)
 
-    @property
-    def required_keys(self) -> set[Tensor]:
-        return self._required_keys
-
-    @property
-    def output_keys(self) -> set[Tensor]:
-        return self._output_keys
+    def check_keys(self) -> tuple[set[Tensor], set[Tensor]]:
+        return self._required_keys, self._output_keys
 
 
 def test_call_checks_keys():
@@ -69,10 +64,10 @@ def test_compose_checks_keys():
     t1 = FakeTransform(required_keys={a1}, output_keys={a1, a2})
     t2 = FakeTransform(required_keys={a2}, output_keys={a1})
 
-    t1 << t2
+    (t1 << t2).check_keys()
 
     with raises(ValueError):
-        t2 << t1
+        (t2 << t1).check_keys()
 
 
 def test_conjunct_checks_required_keys():
@@ -88,13 +83,13 @@ def test_conjunct_checks_required_keys():
     t2 = FakeTransform(required_keys={a1}, output_keys=set())
     t3 = FakeTransform(required_keys={a2}, output_keys=set())
 
-    t1 | t2
+    (t1 | t2).check_keys()
 
     with raises(ValueError):
-        t2 | t3
+        (t2 | t3).check_keys()
 
     with raises(ValueError):
-        t1 | t2 | t3
+        (t1 | t2 | t3).check_keys()
 
 
 def test_conjunct_checks_output_keys():
@@ -110,13 +105,13 @@ def test_conjunct_checks_output_keys():
     t2 = FakeTransform(required_keys=set(), output_keys={a1})
     t3 = FakeTransform(required_keys=set(), output_keys={a2})
 
-    t2 | t3
+    (t2 | t3).check_keys()
 
     with raises(ValueError):
-        t1 | t3
+        (t1 | t3).check_keys()
 
     with raises(ValueError):
-        t1 | t2 | t3
+        (t1 | t2 | t3).check_keys()
 
 
 def test_empty_conjunction():

tests/unit/autojac/_transform/test_interactions.py

@@ -262,4 +262,4 @@ def test_stack_different_required_keys():
     grad2 = Grad([y2], [a])
 
     with raises(ValueError):
-        _ = Stack([grad1, grad2])
+        Stack([grad1, grad2]).check_keys()

tests/unit/autojac/_transform/test_stack.py

@@ -20,13 +20,8 @@ def __init__(self, keys: Iterable[Tensor]):
     def _compute(self, input: EmptyTensorDict) -> Gradients:
         return Gradients({key: torch.ones_like(key) for key in self.keys})
 
-    @property
-    def required_keys(self) -> set[Tensor]:
-        return set()
-
-    @property
-    def output_keys(self) -> set[Tensor]:
-        return self.keys
+    def check_keys(self) -> tuple[set[Tensor], set[Tensor]]:
+        return set(), self.keys
 
 
 def test_single_key():