Make OrderedSet public (#356)

* Import OrderedSet in _transform/__init__.py
* Import OrderedSet directly from _transform rather than from _transform._ordered_set

Author: ValerianRey

src/torchjd/_autojac/_backward.py

@@ -4,8 +4,16 @@
 
 from torchjd.aggregation import Aggregator
 
-from ._transform import Accumulate, Aggregate, Diagonalize, EmptyTensorDict, Init, Jac, Transform
-from ._transform._ordered_set import OrderedSet
+from ._transform import (
+    Accumulate,
+    Aggregate,
+    Diagonalize,
+    EmptyTensorDict,
+    Init,
+    Jac,
+    OrderedSet,
+    Transform,
+)
 from ._utils import as_checked_ordered_set, check_optional_positive_chunk_size, get_leaf_tensors
 
 

src/torchjd/_autojac/_mtl_backward.py

@@ -12,11 +12,11 @@
     Gradients,
     Init,
     Jac,
+    OrderedSet,
     Select,
     Stack,
     Transform,
 )
-from ._transform._ordered_set import OrderedSet
 from ._utils import as_checked_ordered_set, check_optional_positive_chunk_size, get_leaf_tensors
 
 

src/torchjd/_autojac/_transform/__init__.py

@@ -5,6 +5,7 @@
 from ._grad import Grad
 from ._init import Init
 from ._jac import Jac
+from ._ordered_set import OrderedSet
 from ._select import Select
 from ._stack import Stack
 from ._tensor_dict import (

src/torchjd/_autojac/_utils.py

@@ -4,7 +4,7 @@
 from torch import Tensor
 from torch.autograd.graph import Node
 
-from ._transform._ordered_set import OrderedSet
+from ._transform import OrderedSet
 
 
 def check_optional_positive_chunk_size(parallel_chunk_size: int | None) -> None:

tests/unit/autojac/_transform/test_aggregate.py

@@ -10,10 +10,10 @@
     GradientVectors,
     JacobianMatrices,
     Jacobians,
+    OrderedSet,
     RequirementError,
 )
 from torchjd._autojac._transform._aggregate import _AggregateMatrices, _Matrixify, _Reshape
-from torchjd._autojac._transform._ordered_set import OrderedSet
 from torchjd.aggregation import Random
 
 from ._dict_assertions import assert_tensor_dicts_are_close

tests/unit/autojac/_transform/test_diagonalize.py

@@ -1,8 +1,7 @@
 import torch
 from pytest import raises
 
-from torchjd._autojac._transform import Diagonalize, Gradients, RequirementError
-from torchjd._autojac._transform._ordered_set import OrderedSet
+from torchjd._autojac._transform import Diagonalize, Gradients, OrderedSet, RequirementError
 
 from ._dict_assertions import assert_tensor_dicts_are_close
 

tests/unit/autojac/_transform/test_grad.py

@@ -1,8 +1,7 @@
 import torch
 from pytest import raises
 
-from torchjd._autojac._transform import Grad, Gradients, RequirementError
-from torchjd._autojac._transform._ordered_set import OrderedSet
+from torchjd._autojac._transform import Grad, Gradients, OrderedSet, RequirementError
 
 from ._dict_assertions import assert_tensor_dicts_are_close
 

tests/unit/autojac/_transform/test_interactions.py

@@ -12,12 +12,12 @@
     Init,
     Jac,
     Jacobians,
+    OrderedSet,
     RequirementError,
     Select,
     Stack,
     TensorDict,
 )
-from torchjd._autojac._transform._ordered_set import OrderedSet
 
 from ._dict_assertions import assert_tensor_dicts_are_close
 

tests/unit/autojac/_transform/test_jac.py

@@ -1,8 +1,7 @@
 import torch
 from pytest import mark, raises
 
-from torchjd._autojac._transform import Jac, Jacobians, RequirementError
-from torchjd._autojac._transform._ordered_set import OrderedSet
+from torchjd._autojac._transform import Jac, Jacobians, OrderedSet, RequirementError
 
 from ._dict_assertions import assert_tensor_dicts_are_close
 

tests/unit/autojac/test_backward.py

@@ -5,7 +5,7 @@
 
 from torchjd import backward
 from torchjd._autojac._backward import _create_transform
-from torchjd._autojac._transform._ordered_set import OrderedSet
+from torchjd._autojac._transform import OrderedSet
 from torchjd.aggregation import MGDA, Aggregator, Mean, Random, Sum, UPGrad