fix(autojac): Fix non-deterministic column order (#283)

* Replace the ordered_set function and the _OrderedSet TypeVar in _utils.py by the OrderedSet class in ordered_set.py
* Add difference_update and add to OrderedSet
* Change _get_descendant_accumulate_grads and _get_leaf_tensors to work with OrderedSets
* Make backward and mtl_backward use OrderedSet instead of set or list for tensors to differentiate. This should make the _AggregateMatrices transform use a deterministic key_order, which should fix the column order before aggregation.
* Add changelog entry

Note: we did not verify that there was some non-determinism in the column ordering, and we think that with mtl_backward it could only happen when the parameters were not specified by the user. Still, this should make things much safer.

Author: ValerianRey

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.

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 .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
 

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

@@ -6,8 +6,8 @@
 
 from torchjd.aggregation import Aggregator
 
-from ._utils import _OrderedSet, ordered_set
 from .base import Transform
+from .ordered_set import OrderedSet
 from .tensor_dict import EmptyTensorDict, Gradients, GradientVectors, JacobianMatrices, Jacobians
 
 _KeyType = TypeVar("_KeyType", bound=Hashable)
@@ -32,7 +32,7 @@ def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
 
 class _AggregateMatrices(Transform[JacobianMatrices, GradientVectors]):
     def __init__(self, aggregator: Aggregator, key_order: Iterable[Tensor]):
-        self.key_order = ordered_set(key_order)
+        self.key_order = OrderedSet(key_order)
         self.aggregator = aggregator
 
     def __call__(self, jacobian_matrices: JacobianMatrices) -> GradientVectors:
@@ -54,7 +54,7 @@ def check_and_get_keys(self) -> tuple[set[Tensor], set[Tensor]]:
 
     @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``.

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 .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:

src/torchjd/autojac/_transform/ordered_set.py

@@ -0,0 +1,23 @@
+from collections import OrderedDict
+from typing import Iterable
+
+from torchjd.autojac._transform._utils import _KeyType
+
+
+class OrderedSet(OrderedDict[_KeyType, None]):
+    """Ordered collection of distinct elements."""
+
+    def __init__(self, elements: Iterable[_KeyType]):
+        super().__init__([(element, None) for element in elements])
+
+    def difference_update(self, elements: set[_KeyType]) -> None:
+        """Removes all specified elements from the OrderedSet."""
+
+        for element in elements:
+            if element in self:
+                del self[element]
+
+    def add(self, element: _KeyType) -> None:
+        """Adds the specified element to the OrderedSet."""
+
+        self[element] = None

src/torchjd/autojac/_utils.py

@@ -4,6 +4,8 @@
 from torch import Tensor
 from torch.autograd.graph import Node
 
+from ._transform.ordered_set import OrderedSet
+
 
 def _check_optional_positive_chunk_size(parallel_chunk_size: int | None) -> None:
     if not (parallel_chunk_size is None or parallel_chunk_size > 0):
@@ -21,7 +23,7 @@ def _as_tensor_list(tensors: Sequence[Tensor] | Tensor) -> list[Tensor]:
     return output
 
 
-def _get_leaf_tensors(tensors: Iterable[Tensor], excluded: Iterable[Tensor]) -> set[Tensor]:
+def _get_leaf_tensors(tensors: Iterable[Tensor], excluded: Iterable[Tensor]) -> OrderedSet[Tensor]:
     """
     Gets the leaves of the autograd graph of all specified ``tensors``.
 
@@ -39,15 +41,17 @@ def _get_leaf_tensors(tensors: Iterable[Tensor], excluded: Iterable[Tensor]) ->
         raise ValueError("All `excluded` tensors should have a `grad_fn`.")
 
     accumulate_grads = _get_descendant_accumulate_grads(
-        roots={tensor.grad_fn for tensor in tensors},
+        roots=OrderedSet([tensor.grad_fn for tensor in tensors]),
         excluded_nodes={tensor.grad_fn for tensor in excluded},
     )
-    leaves = {g.variable for g in accumulate_grads}
+    leaves = OrderedSet([g.variable for g in accumulate_grads])
 
     return leaves
 
 
-def _get_descendant_accumulate_grads(roots: set[Node], excluded_nodes: set[Node]) -> set[Node]:
+def _get_descendant_accumulate_grads(
+    roots: OrderedSet[Node], excluded_nodes: set[Node]
+) -> OrderedSet[Node]:
     """
     Gets the AccumulateGrad descendants of the specified nodes.
 
@@ -56,8 +60,9 @@ def _get_descendant_accumulate_grads(roots: set[Node], excluded_nodes: set[Node]
     """
 
     excluded_nodes = set(excluded_nodes)  # Re-instantiate set to avoid modifying input
-    result = set()
-    nodes_to_traverse = deque(roots - excluded_nodes)
+    result = OrderedSet([])
+    roots.difference_update(excluded_nodes)
+    nodes_to_traverse = deque(roots)
 
     # This implementation more or less follows what is advised in
     # https://discuss.pytorch.org/t/autograd-graph-traversal/213658 and what was suggested in

src/torchjd/autojac/backward.py

@@ -5,6 +5,7 @@
 from torchjd.aggregation import Aggregator
 
 from ._transform import Accumulate, Aggregate, Diagonalize, EmptyTensorDict, Init, Jac, Transform
+from ._transform.ordered_set import OrderedSet
 from ._utils import _as_tensor_list, _check_optional_positive_chunk_size, _get_leaf_tensors
 
 
@@ -76,7 +77,7 @@ def backward(
     if inputs is None:
         inputs = _get_leaf_tensors(tensors=tensors, excluded=set())
     else:
-        inputs = set(inputs)
+        inputs = OrderedSet(inputs)
 
     backward_transform = _create_transform(
         tensors=tensors,
@@ -92,7 +93,7 @@ def backward(
 def _create_transform(
     tensors: list[Tensor],
     aggregator: Aggregator,
-    inputs: set[Tensor],
+    inputs: OrderedSet[Tensor],
     retain_graph: bool,
     parallel_chunk_size: int | None,
 ) -> Transform[EmptyTensorDict, EmptyTensorDict]:

src/torchjd/autojac/mtl_backward.py

@@ -16,6 +16,7 @@
     Stack,
     Transform,
 )
+from ._transform.ordered_set import OrderedSet
 from ._utils import _as_tensor_list, _check_optional_positive_chunk_size, _get_leaf_tensors
 
 
@@ -84,8 +85,12 @@ def mtl_backward(
 
     if shared_params is None:
         shared_params = _get_leaf_tensors(tensors=features, excluded=[])
+    else:
+        shared_params = OrderedSet(shared_params)
     if tasks_params is None:
         tasks_params = [_get_leaf_tensors(tensors=[loss], excluded=features) for loss in losses]
+    else:
+        tasks_params = [OrderedSet(task_params) for task_params in tasks_params]
 
     if len(features) == 0:
         raise ValueError("`features` cannot be empty.")
@@ -115,8 +120,8 @@ def _create_transform(
     losses: Sequence[Tensor],
     features: list[Tensor],
     aggregator: Aggregator,
-    tasks_params: list[Iterable[Tensor]],
-    shared_params: set[Tensor],
+    tasks_params: list[OrderedSet[Tensor]],
+    shared_params: OrderedSet[Tensor],
     retain_graph: bool,
     parallel_chunk_size: int | None,
 ) -> Transform[EmptyTensorDict, EmptyTensorDict]:
@@ -126,9 +131,6 @@ def _create_transform(
     task-specific parameters).
     """
 
-    shared_params = list(shared_params)
-    tasks_params = [list(task_params) for task_params in tasks_params]
-
     # Task-specific transforms. Each of them computes and accumulates the gradient of the task's
     # loss w.r.t. the task's specific parameters, and computes and backpropagates the gradient of
     # the losses w.r.t. the shared representations.
@@ -160,12 +162,13 @@ def _create_transform(
 
 def _create_task_transform(
     features: list[Tensor],
-    task_params: list[Tensor],
+    task_params: OrderedSet[Tensor],
     loss: Tensor,
     retain_graph: bool,
 ) -> Transform[EmptyTensorDict, Gradients]:
     # Tensors with respect to which we compute the gradients.
-    to_differentiate = task_params + features
+    to_differentiate = OrderedSet(task_params)  # Re-instantiate set to avoid modifying input
+    to_differentiate.update(OrderedSet(features))
 
     # Transform that initializes the gradient output to 1.
     init = Init([loss])

tests/unit/autojac/test_utils.py

@@ -15,7 +15,7 @@ def test_simple_get_leaf_tensors():
     y2 = (a1**2).sum() + a2.norm()
 
     leaves = _get_leaf_tensors(tensors=[y1, y2], excluded=set())
-    assert leaves == {a1, a2}
+    assert set(leaves) == {a1, a2}
 
 
 def test_get_leaf_tensors_excluded_1():
@@ -36,7 +36,7 @@ def test_get_leaf_tensors_excluded_1():
     y2 = b1
 
     leaves = _get_leaf_tensors(tensors=[y1, y2], excluded={b1, b2})
-    assert leaves == {a1}
+    assert set(leaves) == {a1}
 
 
 def test_get_leaf_tensors_excluded_2():
@@ -57,7 +57,7 @@ def test_get_leaf_tensors_excluded_2():
     y2 = b1
 
     leaves = _get_leaf_tensors(tensors=[y1, y2], excluded={b1, b2})
-    assert leaves == {a1, a2}
+    assert set(leaves) == {a1, a2}
 
 
 def test_get_leaf_tensors_leaf_not_requiring_grad():
@@ -72,7 +72,7 @@ def test_get_leaf_tensors_leaf_not_requiring_grad():
     y2 = (a1**2).sum() + a2.norm()
 
     leaves = _get_leaf_tensors(tensors=[y1, y2], excluded=set())
-    assert leaves == {a1}
+    assert set(leaves) == {a1}
 
 
 def test_get_leaf_tensors_model():
@@ -91,7 +91,7 @@ def test_get_leaf_tensors_model():
     losses = loss_fn(y_hat, y)
 
     leaves = _get_leaf_tensors(tensors=[losses], excluded=set())
-    assert leaves == set(model.parameters())
+    assert set(leaves) == set(model.parameters())
 
 
 def test_get_leaf_tensors_model_excluded_2():
@@ -112,7 +112,7 @@ def test_get_leaf_tensors_model_excluded_2():
     losses = loss_fn(z_hat, z)
 
     leaves = _get_leaf_tensors(tensors=[losses], excluded={y})
-    assert leaves == set(model2.parameters())
+    assert set(leaves) == set(model2.parameters())
 
 
 def test_get_leaf_tensors_single_root():
@@ -122,14 +122,14 @@ def test_get_leaf_tensors_single_root():
     y = p * 2
 
     leaves = _get_leaf_tensors(tensors=[y], excluded=set())
-    assert leaves == {p}
+    assert set(leaves) == {p}
 
 
 def test_get_leaf_tensors_empty_roots():
     """Tests that _get_leaf_tensors returns no leaves when roots is the empty set."""
 
     leaves = _get_leaf_tensors(tensors=[], excluded=set())
-    assert leaves == set()
+    assert set(leaves) == set({})
 
 
 def test_get_leaf_tensors_excluded_root():
@@ -142,7 +142,7 @@ def test_get_leaf_tensors_excluded_root():
     y2 = (a1**2).sum()
 
     leaves = _get_leaf_tensors(tensors=[y1, y2], excluded={y1})
-    assert leaves == {a1}
+    assert set(leaves) == {a1}
 
 
 @mark.parametrize("depth", [100, 1000, 10000])
@@ -155,7 +155,7 @@ def test_get_leaf_tensors_deep(depth: int):
         sum_ = sum_ + one
 
     leaves = _get_leaf_tensors(tensors=[sum_], excluded=set())
-    assert leaves == {one}
+    assert set(leaves) == {one}
 
 
 def test_get_leaf_tensors_leaf():