Make more things protected

CHANGELOG.md

@@ -14,6 +14,13 @@ changes that do not affect the user.
   TorchJD. Users of these aggregators will have to use `pip install torchjd[cagrad]`, `pip install
   torchjd[nash_mtl]` or `pip install torchjd[full]` to install TorchJD alongside those dependencies.
   This should make TorchJD more lightweight.
+- **BREAKING**: Made the aggregator modules and the `autojac` package protected. The aggregators
+  must now always be imported via their package (e.g.
+  `from torchjd.aggregation.upgrad import UPGrad` must be changed to
+  `from torchjd.aggregation import UPGrad`). The `backward` and `mtl_backward` functions must now
+  always be imported directly from the `torchjd` package (e.g.
+  `from torchjd.autojac.mtl_backward import mtl_backward` must be changed to
+  `from torchjd import mtl_backward`).
 
 ### Fixed
 

docs/source/docs/aggregation/aligned_mtl.rst

@@ -3,7 +3,7 @@
 Aligned-MTL
 ===========
 
-.. automodule:: torchjd.aggregation.aligned_mtl
+.. autoclass:: torchjd.aggregation.AlignedMTL
     :members:
     :undoc-members:
     :exclude-members: forward

docs/source/docs/aggregation/bases.rst

@@ -3,7 +3,7 @@
 Aggregator (abstract)
 =====================
 
-.. automodule:: torchjd.aggregation.aggregator_bases
+.. autoclass:: torchjd.aggregation.Aggregator
     :members:
     :undoc-members:
     :show-inheritance:

docs/source/docs/aggregation/cagrad.rst

@@ -3,7 +3,7 @@
 CAGrad
 ======
 
-.. automodule:: torchjd.aggregation.cagrad
+.. autoclass:: torchjd.aggregation.CAGrad
     :members:
     :undoc-members:
     :exclude-members: forward

docs/source/docs/aggregation/config.rst

@@ -3,7 +3,7 @@
 ConFIG
 ======
 
-.. automodule:: torchjd.aggregation.config
+.. autoclass:: torchjd.aggregation.ConFIG
     :members:
     :undoc-members:
     :exclude-members: forward

docs/source/docs/aggregation/constant.rst

@@ -3,7 +3,7 @@
 Constant
 ========
 
-.. automodule:: torchjd.aggregation.constant
+.. autoclass:: torchjd.aggregation.Constant
     :members:
     :undoc-members:
     :exclude-members: forward

docs/source/docs/aggregation/dualproj.rst

@@ -3,7 +3,7 @@
 DualProj
 ========
 
-.. automodule:: torchjd.aggregation.dualproj
+.. autoclass:: torchjd.aggregation.DualProj
     :members:
     :undoc-members:
     :exclude-members: forward

docs/source/docs/aggregation/graddrop.rst

@@ -3,7 +3,7 @@
 GradDrop
 ========
 
-.. automodule:: torchjd.aggregation.graddrop
+.. autoclass:: torchjd.aggregation.GradDrop
     :members:
     :undoc-members:
     :exclude-members: forward

docs/source/docs/aggregation/imtl_g.rst

@@ -3,7 +3,7 @@
 IMTL-G
 ======
 
-.. automodule:: torchjd.aggregation.imtl_g
+.. autoclass:: torchjd.aggregation.IMTLG
     :members:
     :undoc-members:
     :exclude-members: forward

docs/source/docs/aggregation/krum.rst

@@ -3,7 +3,7 @@
 Krum
 ====
 
-.. automodule:: torchjd.aggregation.krum
+.. autoclass:: torchjd.aggregation.Krum
     :members:
     :undoc-members:
     :exclude-members: forward

docs/source/docs/aggregation/mean.rst

@@ -3,7 +3,7 @@
 Mean
 ====
 
-.. automodule:: torchjd.aggregation.mean
+.. autoclass:: torchjd.aggregation.Mean
     :members:
     :undoc-members:
     :exclude-members: forward

docs/source/docs/aggregation/mgda.rst

@@ -3,7 +3,7 @@
 MGDA
 ====
 
-.. automodule:: torchjd.aggregation.mgda
+.. autoclass:: torchjd.aggregation.MGDA
     :members:
     :undoc-members:
     :exclude-members: forward

docs/source/docs/aggregation/nash_mtl.rst

@@ -3,7 +3,7 @@
 Nash-MTL
 ========
 
-.. automodule:: torchjd.aggregation.nash_mtl
+.. autoclass:: torchjd.aggregation.NashMTL
     :members:
     :undoc-members:
     :exclude-members: forward

docs/source/docs/aggregation/pcgrad.rst

@@ -3,7 +3,7 @@
 PCGrad
 ======
 
-.. automodule:: torchjd.aggregation.pcgrad
+.. autoclass:: torchjd.aggregation.PCGrad
     :members:
     :undoc-members:
     :exclude-members: forward

docs/source/docs/aggregation/random.rst

@@ -3,7 +3,7 @@
 Random
 ======
 
-.. automodule:: torchjd.aggregation.random
+.. autoclass:: torchjd.aggregation.Random
     :members:
     :undoc-members:
     :exclude-members: forward

docs/source/docs/aggregation/sum.rst

@@ -3,7 +3,7 @@
 Sum
 ===
 
-.. automodule:: torchjd.aggregation.sum
+.. autoclass:: torchjd.aggregation.Sum
     :members:
     :undoc-members:
     :exclude-members: forward

docs/source/docs/aggregation/trimmed_mean.rst

@@ -3,7 +3,7 @@
 Trimmed Mean
 ============
 
-.. automodule:: torchjd.aggregation.trimmed_mean
+.. autoclass:: torchjd.aggregation.TrimmedMean
     :members:
     :undoc-members:
     :exclude-members: forward

docs/source/docs/aggregation/upgrad.rst

@@ -3,7 +3,7 @@
 UPGrad
 ======
 
-.. automodule:: torchjd.aggregation.upgrad
+.. autoclass:: torchjd.aggregation.UPGrad
     :members:
     :undoc-members:
     :exclude-members: forward

docs/source/docs/autojac/backward.rst

@@ -3,7 +3,4 @@
 backward
 ========
 
-.. automodule:: torchjd.autojac.backward
-    :members:
-    :undoc-members:
-    :exclude-members:
+.. autofunction:: torchjd.backward

docs/source/docs/autojac/mtl_backward.rst

@@ -3,7 +3,4 @@
 mtl_backward
 ============
 
-.. automodule:: torchjd.autojac.mtl_backward
-    :members:
-    :undoc-members:
-    :exclude-members:
+.. autofunction:: torchjd.mtl_backward

src/torchjd/__init__.py

@@ -1 +1 @@
-from torchjd.autojac import backward, mtl_backward
+from ._autojac import backward, mtl_backward

src/torchjd/_autojac/__init__.py

@@ -0,0 +1,2 @@
+from ._backward import backward
+from ._mtl_backward import mtl_backward

src/torchjd/autojac/backward.py → src/torchjd/_autojac/_backward.py

@@ -5,7 +5,7 @@
 from torchjd.aggregation import Aggregator
 
 from ._transform import Accumulate, Aggregate, Diagonalize, EmptyTensorDict, Init, Jac, Transform
-from ._transform.ordered_set import OrderedSet
+from ._transform._ordered_set import OrderedSet
 from ._utils import as_checked_ordered_set, check_optional_positive_chunk_size, get_leaf_tensors
 
 

src/torchjd/autojac/mtl_backward.py → src/torchjd/_autojac/_mtl_backward.py

@@ -16,7 +16,7 @@
     Stack,
     Transform,
 )
-from ._transform.ordered_set import OrderedSet
+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

@@ -0,0 +1,17 @@
+from ._accumulate import Accumulate
+from ._aggregate import Aggregate
+from ._base import Composition, Conjunction, RequirementError, Transform
+from ._diagonalize import Diagonalize
+from ._grad import Grad
+from ._init import Init
+from ._jac import Jac
+from ._select import Select
+from ._stack import Stack
+from ._tensor_dict import (
+    EmptyTensorDict,
+    Gradients,
+    GradientVectors,
+    JacobianMatrices,
+    Jacobians,
+    TensorDict,
+)

src/torchjd/autojac/_transform/accumulate.py → src/torchjd/_autojac/_transform/_accumulate.py

@@ -1,7 +1,7 @@
 from torch import Tensor
 
-from .base import Transform
-from .tensor_dict import EmptyTensorDict, Gradients
+from ._base import Transform
+from ._tensor_dict import EmptyTensorDict, Gradients
 
 
 class Accumulate(Transform[Gradients, EmptyTensorDict]):

src/torchjd/autojac/_transform/aggregate.py → src/torchjd/_autojac/_transform/_aggregate.py

@@ -6,9 +6,9 @@
 
 from torchjd.aggregation import Aggregator
 
-from .base import RequirementError, Transform
-from .ordered_set import OrderedSet
-from .tensor_dict import EmptyTensorDict, Gradients, GradientVectors, JacobianMatrices, Jacobians
+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")
@@ -94,7 +94,7 @@ def _aggregate_group(
     ) -> GradientVectors:
         """
         Unites the jacobian matrices and aggregates them using an
-        :class:`~torchjd.aggregation.aggregator_bases.Aggregator`. Returns the obtained gradient
+        :class:`~torchjd.aggregation._aggregator_bases.Aggregator`. Returns the obtained gradient
         vectors.
         """
 

src/torchjd/autojac/_transform/base.py → src/torchjd/_autojac/_transform/_base.py

@@ -5,7 +5,7 @@
 
 from torch import Tensor
 
-from .tensor_dict import _A, _B, _C, EmptyTensorDict, _least_common_ancestor
+from ._tensor_dict import _A, _B, _C, EmptyTensorDict, _least_common_ancestor
 
 
 class RequirementError(ValueError):

src/torchjd/autojac/_transform/diagonalize.py → src/torchjd/_autojac/_transform/_diagonalize.py

@@ -1,9 +1,9 @@
 import torch
 from torch import Tensor
 
-from .base import RequirementError, Transform
-from .ordered_set import OrderedSet
-from .tensor_dict import Gradients, Jacobians
+from ._base import RequirementError, Transform
+from ._ordered_set import OrderedSet
+from ._tensor_dict import Gradients, Jacobians
 
 
 class Diagonalize(Transform[Gradients, Jacobians]):

src/torchjd/autojac/_transform/_differentiate.py → src/torchjd/_autojac/_transform/_differentiate.py

@@ -3,9 +3,9 @@
 
 from torch import Tensor
 
-from .base import RequirementError, Transform
-from .ordered_set import OrderedSet
-from .tensor_dict import _A
+from ._base import RequirementError, Transform
+from ._ordered_set import OrderedSet
+from ._tensor_dict import _A
 
 
 class Differentiate(Transform[_A, _A], ABC):

src/torchjd/autojac/_transform/grad.py → src/torchjd/_autojac/_transform/_grad.py

@@ -5,8 +5,8 @@
 
 from ._differentiate import Differentiate
 from ._materialize import materialize
-from .ordered_set import OrderedSet
-from .tensor_dict import Gradients
+from ._ordered_set import OrderedSet
+from ._tensor_dict import Gradients
 
 
 class Grad(Differentiate[Gradients]):

src/torchjd/autojac/_transform/init.py → src/torchjd/_autojac/_transform/_init.py

@@ -3,8 +3,8 @@
 import torch
 from torch import Tensor
 
-from .base import RequirementError, Transform
-from .tensor_dict import EmptyTensorDict, Gradients
+from ._base import RequirementError, Transform
+from ._tensor_dict import EmptyTensorDict, Gradients
 
 
 class Init(Transform[EmptyTensorDict, Gradients]):

src/torchjd/autojac/_transform/jac.py → src/torchjd/_autojac/_transform/_jac.py

@@ -8,8 +8,8 @@
 
 from ._differentiate import Differentiate
 from ._materialize import materialize
-from .ordered_set import OrderedSet
-from .tensor_dict import Jacobians
+from ._ordered_set import OrderedSet
+from ._tensor_dict import Jacobians
 
 
 class Jac(Differentiate[Jacobians]):

src/torchjd/autojac/_transform/select.py → src/torchjd/_autojac/_transform/_select.py

@@ -2,8 +2,8 @@
 
 from torch import Tensor
 
-from .base import RequirementError, Transform
-from .tensor_dict import _A
+from ._base import RequirementError, Transform
+from ._tensor_dict import _A
 
 
 class Select(Transform[_A, _A]):

src/torchjd/autojac/_transform/stack.py → src/torchjd/_autojac/_transform/_stack.py

@@ -3,9 +3,9 @@
 import torch
 from torch import Tensor
 
+from ._base import Transform
 from ._materialize import materialize
-from .base import Transform
-from .tensor_dict import _A, Gradients, Jacobians
+from ._tensor_dict import _A, Gradients, Jacobians
 
 
 class Stack(Transform[_A, Jacobians]):

src/torchjd/autojac/_utils.py → 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._ordered_set import OrderedSet
 
 
 def check_optional_positive_chunk_size(parallel_chunk_size: int | None) -> None: