Merge branch 'main' into improve-aggregate-typing

Author: ValerianRey

tests/unit/autojac/test_backward.py

@@ -5,6 +5,7 @@
 
 from torchjd import backward
 from torchjd.aggregation import MGDA, Aggregator, Mean, Random, Sum, UPGrad
+from torchjd.autojac._transform.ordered_set import OrderedSet
 from torchjd.autojac.backward import _create_transform
 
 
@@ -20,7 +21,7 @@ def test_check_create_transform():
     transform = _create_transform(
         tensors=[y1, y2],
         aggregator=Mean(),
-        inputs={a1, a2},
+        inputs=OrderedSet([a1, a2]),
         retain_graph=False,
         parallel_chunk_size=None,
     )

tests/unit/autojac/test_mtl_backward.py

@@ -5,6 +5,7 @@
 
 from torchjd import mtl_backward
 from torchjd.aggregation import MGDA, Aggregator, Mean, Random, Sum, UPGrad
+from torchjd.autojac._transform.ordered_set import OrderedSet
 from torchjd.autojac.mtl_backward import _create_transform
 
 
@@ -24,8 +25,8 @@ def test_check_create_transform():
         losses=[y1, y2],
         features=[f1, f2],
         aggregator=Mean(),
-        tasks_params=[[p1], [p2]],
-        shared_params={p0},
+        tasks_params=[OrderedSet([p1]), OrderedSet([p2])],
+        shared_params=OrderedSet([p0]),
         retain_graph=False,
         parallel_chunk_size=None,
     )