Merge branch 'main' into revamp-config

Author: ValerianRey

.pre-commit-config.yaml

@@ -10,7 +10,7 @@ repos:
     -   id: check-merge-conflict  # Check for files that contain merge conflict strings.
 
 -   repo: https://github.com/PyCQA/flake8
-    rev: 7.1.2
+    rev: 7.2.0
     hooks:
     -   id: flake8  # Check style and syntax. Does not modify code, issues have to be solved manually.
         args: [

CHANGELOG.md

@@ -21,11 +21,15 @@ changes that do not affect the user.
   and `mtl_backward`.
 
 ### Fixed
+
 - Fixed the behavior of `backward` and `mtl_backward` when some tensors are repeated (i.e. when they
   appear several times in a list of tensors provided as argument). Instead of raising an exception
   in these cases, we are now aligned with the behavior of `torch.autograd.backward`. Repeated
   tensors that we differentiate lead to repeated rows in the Jacobian, prior to aggregation, and
   repeated tensors with respect to which we differentiate count only once.
+- Fixed an issue with `backward` and `mtl_backward` that could make the ordering of the columns of
+  the Jacobians non-deterministic, and that could thus lead to slightly non-deterministic results
+  with some aggregators.
 - Removed arbitrary exception handling in `IMTLG` and `AlignedMTL` when the computation fails. In
   practice, this fix should only affect some matrices with extremely large values, which should
   not usually happen.

README.md

@@ -6,6 +6,7 @@
 [![pre-commit.ci status](https://results.pre-commit.ci/badge/github/TorchJD/torchjd/main.svg)](https://results.pre-commit.ci/latest/github/TorchJD/torchjd/main)
 [![PyPI - Downloads](https://img.shields.io/pypi/dm/torchjd)](https://pypistats.org/packages/torchjd)
 [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/torchjd)](https://pypi.org/project/torchjd/)
+[![Static Badge](https://img.shields.io/badge/Discord%20-%20community%20-%20%235865F2?logo=discord&logoColor=%23FFFFFF&label=Discord)](https://discord.gg/76KkRnb3nk)
 
 TorchJD is a library extending autograd to enable
 [Jacobian descent](https://arxiv.org/pdf/2406.16232) with PyTorch. It can be used to train neural

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/_differentiate.py

@@ -3,8 +3,8 @@
 
 from torch import Tensor
 
-from ._utils import ordered_set
-from .base import _A, Transform
+from .base import _A, RequirementError, Transform
+from .ordered_set import OrderedSet
 
 
 class _Differentiate(Transform[_A, _A], ABC):
@@ -16,7 +16,7 @@ def __init__(
         create_graph: bool,
     ):
         self.outputs = list(outputs)
-        self.inputs = ordered_set(inputs)
+        self.inputs = OrderedSet(inputs)
         self.retain_graph = retain_graph
         self.create_graph = create_graph
 
@@ -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 == input_keys:
+            raise RequirementError(
+                f"The input_keys must match the expected outputs. Found input_keys {input_keys} and"
+                f"outputs {outputs}."
+            )
+        return set(self.inputs)

src/torchjd/autojac/_transform/_utils.py

@@ -1,5 +1,4 @@
-from collections import OrderedDict
-from typing import Hashable, Iterable, Sequence, TypeAlias, TypeVar
+from typing import Hashable, Iterable, Sequence, TypeVar
 
 import torch
 from torch import Tensor
@@ -8,17 +7,12 @@
 
 _KeyType = TypeVar("_KeyType", bound=Hashable)
 _ValueType = TypeVar("_ValueType")
-_OrderedSet: TypeAlias = OrderedDict[_KeyType, None]
 
 _A = TypeVar("_A", bound=TensorDict)
 _B = TypeVar("_B", bound=TensorDict)
 _C = TypeVar("_C", bound=TensorDict)
 
 
-def ordered_set(elements: Iterable[_KeyType]) -> _OrderedSet[_KeyType]:
-    return OrderedDict.fromkeys(elements, None)
-
-
 def dicts_union(dicts: Iterable[dict[_KeyType, _ValueType]]) -> dict[_KeyType, _ValueType]:
     result = {}
     for d in dicts:

src/torchjd/autojac/_transform/accumulate.py

@@ -1,15 +1,10 @@
-from typing import Iterable
-
 from torch import Tensor
 
 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.
@@ -28,8 +23,8 @@ 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]:
+        return set()
 
 
 def _check_expects_grad(tensor: Tensor) -> None:

src/torchjd/autojac/_transform/aggregate.py

@@ -1,38 +1,38 @@
 from collections import OrderedDict
-from typing import Hashable, Iterable, TypeVar
+from typing import Hashable, TypeVar
 
 import torch
 from torch import Tensor
 
 from torchjd.aggregation import Aggregator
 
-from ._utils import _OrderedSet, ordered_set
-from .base import Transform
+from .base import RequirementError, Transform
+from .ordered_set import OrderedSet
 from .tensor_dict import EmptyTensorDict, Gradients, GradientVectors, JacobianMatrices, Jacobians
 
 _KeyType = TypeVar("_KeyType", bound=Hashable)
 _ValueType = TypeVar("_ValueType")
 
 
 class Aggregate(Transform[Jacobians, Gradients]):
-    def __init__(self, aggregator: Aggregator, key_order: Iterable[Tensor]):
-        matrixify = _Matrixify(key_order)
+    def __init__(self, aggregator: Aggregator, key_order: OrderedSet[Tensor]):
+        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]):
-    def __init__(self, aggregator: Aggregator, key_order: Iterable[Tensor]):
-        self.key_order = ordered_set(key_order)
+    def __init__(self, aggregator: Aggregator, key_order: OrderedSet[Tensor]):
+        self.key_order = OrderedSet(key_order)
         self.aggregator = aggregator
 
     def __call__(self, jacobian_matrices: JacobianMatrices) -> GradientVectors:
@@ -48,13 +48,17 @@ 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 {input_keys} and"
+                f"key_order {self.key_order}."
+            )
+        return input_keys
 
     @staticmethod
     def _select_ordered_subdict(
-        dictionary: dict[_KeyType, _ValueType], ordered_keys: _OrderedSet[_KeyType]
+        dictionary: dict[_KeyType, _ValueType], ordered_keys: OrderedSet[_KeyType]
     ) -> OrderedDict[_KeyType, _ValueType]:
         """
         Selects a subset of a dictionary corresponding to the keys given by ``ordered_keys``.
@@ -108,29 +112,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

@@ -8,6 +8,12 @@
 from ._utils import _A, _B, _C, _union
 
 
+class RequirementError(ValueError):
+    """Inappropriate set of inputs keys."""
+
+    pass
+
+
 class Transform(Generic[_B, _C], ABC):
     r"""
     Abstract base class for all transforms. Transforms are elementary building blocks of a jacobian
@@ -44,17 +50,16 @@ 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 that the provided input_keys satisfy the transform's requirements 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.
+        If the provided input_keys do not satisfy the transform's requirements, raises a
+        RequirementError.
 
-        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 of the output TensorDict of the transform when the input
+        TensorDict's keys are input_keys.
         """
 
     __lshift__ = compose
@@ -73,15 +78,10 @@ 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 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
+    def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
+        intermediate_keys = self.inner.check_keys(input_keys)
+        output_keys = self.outer.check_keys(intermediate_keys)
+        return output_keys
 
 
 class Conjunction(Transform[_A, _B]):
@@ -102,18 +102,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 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)
+    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):
-            raise ValueError("The sets of output keys of transforms should be disjoint.")
+        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,14 +3,14 @@
 import torch
 from torch import Tensor
 
-from ._utils import ordered_set
-from .base import Transform
+from .base import RequirementError, Transform
+from .ordered_set import OrderedSet
 from .tensor_dict import Gradients, Jacobians
 
 
 class Diagonalize(Transform[Gradients, Jacobians]):
     def __init__(self, considered: Iterable[Tensor]):
-        self.considered = ordered_set(considered)
+        self.considered = OrderedSet(considered)
         self.indices: list[tuple[int, int]] = []
         begin = 0
         for tensor in self.considered:
@@ -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 == input_keys:
+            raise RequirementError(
+                f"The input_keys must match the considered keys. Found input_keys {input_keys} and"
+                f"considered keys {considered}."
+            )
+        return considered