refactor(autojac): Add OrderedSet.__add__ (#313)

* Use it in mtl_backward rather than creating a new OrderedSet and updating it

Author: ValerianRey

src/torchjd/autojac/_transform/ordered_set.py

@@ -22,3 +22,8 @@ def add(self, element: _KeyType) -> None:
         """Adds the specified element to the OrderedSet."""
 
         self[element] = None
+
+    def __add__(self, other: "OrderedSet[_KeyType]") -> "OrderedSet[_KeyType]":
+        """Creates a new OrderedSet with the elements of self followed by the elements of other."""
+
+        return OrderedSet([*self, *other])

src/torchjd/autojac/mtl_backward.py

@@ -168,8 +168,7 @@ def _create_task_transform(
     retain_graph: bool,
 ) -> Transform[EmptyTensorDict, Gradients]:
     # Tensors with respect to which we compute the gradients.
-    to_differentiate = OrderedSet(task_params)  # Re-instantiate set to avoid modifying input
-    to_differentiate.update(features)
+    to_differentiate = task_params + features
 
     # Transform that initializes the gradient output to 1.
     init = Init(loss)