Add RequirementError

Author: PierreQuinton

src/torchjd/autojac/_transform/__init__.py

@@ -1,6 +1,6 @@
 from .accumulate import Accumulate
 from .aggregate import Aggregate
-from .base import Composition, Conjunction, Transform
+from .base import Composition, Conjunction, RequirementError, Transform
 from .diagonalize import Diagonalize
 from .grad import Grad
 from .init import Init

src/torchjd/autojac/_transform/accumulate.py

@@ -2,7 +2,7 @@
 
 from torch import Tensor
 
-from .base import Transform
+from .base import RequirementError, Transform
 from .tensor_dict import EmptyTensorDict, Gradients
 
 
@@ -34,7 +34,7 @@ def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
 
 def _check_expects_grad(tensor: Tensor) -> None:
     if not _expects_grad(tensor):
-        raise ValueError(
+        raise RequirementError(
             "Cannot populate the .grad field of a Tensor that does not satisfy:"
             "`tensor.requires_grad and (tensor.is_leaf or tensor.retains_grad)`."
         )

src/torchjd/autojac/_transform/base.py

@@ -8,6 +8,10 @@
 from ._utils import _A, _B, _C, _union
 
 
+class RequirementError(ValueError):
+    pass
+
+
 class Transform(Generic[_B, _C], ABC):
     r"""
     Abstract base class for all transforms. Transforms are elementary building blocks of a jacobian
@@ -77,7 +81,7 @@ 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 ValueError(
+            raise RequirementError(
                 "The `output_keys` of `inner` must match with the `required_keys` of "
                 f"outer. Found {outer_required_keys} and {inner_output_keys}"
             )
@@ -108,12 +112,12 @@ def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
         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.")
+                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)
 
         if len(output_keys) != len(output_keys_with_duplicates):
-            raise ValueError("The sets of output keys of transforms should be disjoint.")
+            raise RequirementError("The sets of output keys of transforms should be disjoint.")
 
         return required_keys, output_keys

src/torchjd/autojac/_transform/select.py

@@ -3,7 +3,7 @@
 from torch import Tensor
 
 from ._utils import _A
-from .base import Transform
+from .base import RequirementError, Transform
 
 
 class Select(Transform[_A, _A]):
@@ -18,6 +18,8 @@ def __call__(self, tensor_dict: _A) -> _A:
     def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
         required_keys = self._required_keys
         if not self.keys.issubset(required_keys):
-            raise ValueError("Parameter `keys` should be a subset of parameter `required_keys`")
+            raise RequirementError(
+                "Parameter `keys` should be a subset of parameter `required_keys`"
+            )
 
         return required_keys, 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 Transform
+from .base import RequirementError, Transform
 from .tensor_dict import Gradients, Jacobians
 
 
@@ -25,7 +25,7 @@ def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
 
         for transform_required_keys, _ in keys_pairs:
             if transform_required_keys != required_keys:
-                raise ValueError("All transforms should require the same set of keys.")
+                raise RequirementError("All transforms should require the same set of keys.")
 
         return required_keys, output_keys