docs(aggregation): Document the __call__ method publicly (#654)

ValerianRey · web-flow · commit 4597af8cd079 · 2026-04-19T14:02:42.000+02:00
diff --git a/docs/source/docs/aggregation/aligned_mtl.rst b/docs/source/docs/aggregation/aligned_mtl.rst
@@ -4,5 +4,7 @@ Aligned-MTL
 ===========
 
 .. autoclass:: torchjd.aggregation.AlignedMTL
+    :members: __call__
 
 .. autoclass:: torchjd.aggregation.AlignedMTLWeighting
+    :members: __call__
diff --git a/docs/source/docs/aggregation/cagrad.rst b/docs/source/docs/aggregation/cagrad.rst
@@ -4,5 +4,7 @@ CAGrad
 ======
 
 .. autoclass:: torchjd.aggregation.CAGrad
+    :members: __call__
 
 .. autoclass:: torchjd.aggregation.CAGradWeighting
+    :members: __call__
diff --git a/docs/source/docs/aggregation/config.rst b/docs/source/docs/aggregation/config.rst
@@ -4,3 +4,4 @@ ConFIG
 ======
 
 .. autoclass:: torchjd.aggregation.ConFIG
+    :members: __call__
diff --git a/docs/source/docs/aggregation/constant.rst b/docs/source/docs/aggregation/constant.rst
@@ -4,5 +4,7 @@ Constant
 ========
 
 .. autoclass:: torchjd.aggregation.Constant
+    :members: __call__
 
 .. autoclass:: torchjd.aggregation.ConstantWeighting
+    :members: __call__
diff --git a/docs/source/docs/aggregation/dualproj.rst b/docs/source/docs/aggregation/dualproj.rst
@@ -4,5 +4,7 @@ DualProj
 ========
 
 .. autoclass:: torchjd.aggregation.DualProj
+    :members: __call__
 
 .. autoclass:: torchjd.aggregation.DualProjWeighting
+    :members: __call__
diff --git a/docs/source/docs/aggregation/flattening.rst b/docs/source/docs/aggregation/flattening.rst
@@ -4,3 +4,4 @@ Flattening
 ==========
 
 .. autoclass:: torchjd.aggregation.Flattening
+    :members: __call__
diff --git a/docs/source/docs/aggregation/graddrop.rst b/docs/source/docs/aggregation/graddrop.rst
@@ -4,3 +4,4 @@ GradDrop
 ========
 
 .. autoclass:: torchjd.aggregation.GradDrop
+    :members: __call__
diff --git a/docs/source/docs/aggregation/gradvac.rst b/docs/source/docs/aggregation/gradvac.rst
@@ -4,7 +4,7 @@ GradVac
 =======
 
 .. autoclass:: torchjd.aggregation.GradVac
-    :members: reset
+    :members: __call__, reset
 
 .. autoclass:: torchjd.aggregation.GradVacWeighting
-    :members: reset
+    :members: __call__, reset
diff --git a/docs/source/docs/aggregation/imtl_g.rst b/docs/source/docs/aggregation/imtl_g.rst
@@ -4,5 +4,7 @@ IMTL-G
 ======
 
 .. autoclass:: torchjd.aggregation.IMTLG
+    :members: __call__
 
 .. autoclass:: torchjd.aggregation.IMTLGWeighting
+    :members: __call__
diff --git a/docs/source/docs/aggregation/index.rst b/docs/source/docs/aggregation/index.rst
@@ -8,10 +8,13 @@ Abstract base classes
 ---------------------
 
 .. autoclass:: torchjd.aggregation.Aggregator
+    :members: __call__
 
 .. autoclass:: torchjd.aggregation.Weighting
+    :members: __call__
 
 .. autoclass:: torchjd.aggregation.GeneralizedWeighting
+    :members: __call__
 
 .. autoclass:: torchjd.aggregation.Stateful
     :members: reset
diff --git a/docs/source/docs/aggregation/krum.rst b/docs/source/docs/aggregation/krum.rst
@@ -4,5 +4,7 @@ Krum
 ====
 
 .. autoclass:: torchjd.aggregation.Krum
+    :members: __call__
 
 .. autoclass:: torchjd.aggregation.KrumWeighting
+    :members: __call__
diff --git a/docs/source/docs/aggregation/mean.rst b/docs/source/docs/aggregation/mean.rst
@@ -4,5 +4,7 @@ Mean
 ====
 
 .. autoclass:: torchjd.aggregation.Mean
+    :members: __call__
 
 .. autoclass:: torchjd.aggregation.MeanWeighting
+    :members: __call__
diff --git a/docs/source/docs/aggregation/mgda.rst b/docs/source/docs/aggregation/mgda.rst
@@ -4,5 +4,7 @@ MGDA
 ====
 
 .. autoclass:: torchjd.aggregation.MGDA
+    :members: __call__
 
 .. autoclass:: torchjd.aggregation.MGDAWeighting
+    :members: __call__
diff --git a/docs/source/docs/aggregation/nash_mtl.rst b/docs/source/docs/aggregation/nash_mtl.rst
@@ -4,4 +4,4 @@ Nash-MTL
 ========
 
 .. autoclass:: torchjd.aggregation.NashMTL
-    :members: reset
+    :members: __call__, reset
diff --git a/docs/source/docs/aggregation/pcgrad.rst b/docs/source/docs/aggregation/pcgrad.rst
@@ -4,5 +4,7 @@ PCGrad
 ======
 
 .. autoclass:: torchjd.aggregation.PCGrad
+    :members: __call__
 
 .. autoclass:: torchjd.aggregation.PCGradWeighting
+    :members: __call__
diff --git a/docs/source/docs/aggregation/random.rst b/docs/source/docs/aggregation/random.rst
@@ -4,5 +4,7 @@ Random
 ======
 
 .. autoclass:: torchjd.aggregation.Random
+    :members: __call__
 
 .. autoclass:: torchjd.aggregation.RandomWeighting
+    :members: __call__
diff --git a/docs/source/docs/aggregation/sum.rst b/docs/source/docs/aggregation/sum.rst
@@ -4,5 +4,7 @@ Sum
 ===
 
 .. autoclass:: torchjd.aggregation.Sum
+    :members: __call__
 
 .. autoclass:: torchjd.aggregation.SumWeighting
+    :members: __call__
diff --git a/docs/source/docs/aggregation/trimmed_mean.rst b/docs/source/docs/aggregation/trimmed_mean.rst
@@ -4,3 +4,4 @@ Trimmed Mean
 ============
 
 .. autoclass:: torchjd.aggregation.TrimmedMean
+    :members: __call__
diff --git a/docs/source/docs/aggregation/upgrad.rst b/docs/source/docs/aggregation/upgrad.rst
@@ -4,5 +4,7 @@ UPGrad
 ======
 
 .. autoclass:: torchjd.aggregation.UPGrad
+    :members: __call__
 
 .. autoclass:: torchjd.aggregation.UPGradWeighting
+    :members: __call__
diff --git a/src/torchjd/aggregation/_aggregator_bases.py b/src/torchjd/aggregation/_aggregator_bases.py
@@ -29,7 +29,11 @@ def forward(self, matrix: Matrix, /) -> Tensor:
         """Computes the aggregation from the input matrix."""
 
     def __call__(self, matrix: Tensor, /) -> Tensor:
-        """Computes the aggregation from the input matrix and applies all registered hooks."""
+        """
+        Computes the aggregation from the input matrix and applies all registered hooks.
+
+        :param matrix: The Jacobian to aggregate.
+        """
         Aggregator._check_is_matrix(matrix)
         return super().__call__(matrix)
 
diff --git a/src/torchjd/aggregation/_aligned_mtl.py b/src/torchjd/aggregation/_aligned_mtl.py
@@ -12,12 +12,12 @@
 from ._aggregator_bases import GramianWeightedAggregator
 from ._mean import MeanWeighting
 from ._utils.pref_vector import pref_vector_to_str_suffix, pref_vector_to_weighting
-from ._weighting_bases import Weighting
+from ._weighting_bases import GramianWeighting
 
 SUPPORTED_SCALE_MODE: TypeAlias = Literal["min", "median", "rmse"]
 
 
-class AlignedMTLWeighting(Weighting[PSDMatrix]):
+class AlignedMTLWeighting(GramianWeighting):
     r"""
     :class:`~torchjd.aggregation._weighting_bases.Weighting` giving the weights of
     :class:`~torchjd.aggregation.AlignedMTL`.
diff --git a/src/torchjd/aggregation/_cagrad.py b/src/torchjd/aggregation/_cagrad.py
@@ -3,7 +3,7 @@
 from torchjd._linalg import PSDMatrix
 
 from ._utils.check_dependencies import check_dependencies_are_installed
-from ._weighting_bases import Weighting
+from ._weighting_bases import GramianWeighting
 
 check_dependencies_are_installed(["cvxpy", "clarabel"])
 
@@ -18,7 +18,7 @@
 from ._utils.non_differentiable import raise_non_differentiable_error
 
 
-class CAGradWeighting(Weighting[PSDMatrix]):
+class CAGradWeighting(GramianWeighting):
     """
     :class:`~torchjd.aggregation._weighting_bases.Weighting` giving the weights of
     :class:`~torchjd.aggregation.CAGrad`.
diff --git a/src/torchjd/aggregation/_constant.py b/src/torchjd/aggregation/_constant.py
@@ -1,13 +1,11 @@
 from torch import Tensor
 
-from torchjd._linalg import Matrix
-
 from ._aggregator_bases import WeightedAggregator
 from ._utils.str import vector_to_str
-from ._weighting_bases import Weighting
+from ._weighting_bases import MatrixWeighting
 
 
-class ConstantWeighting(Weighting[Matrix]):
+class ConstantWeighting(MatrixWeighting):
     """
     :class:`~torchjd.aggregation._weighting_bases.Weighting` that returns constant, pre-determined
     weights.
diff --git a/src/torchjd/aggregation/_dualproj.py b/src/torchjd/aggregation/_dualproj.py
@@ -7,10 +7,10 @@
 from ._utils.dual_cone import SUPPORTED_SOLVER, project_weights
 from ._utils.non_differentiable import raise_non_differentiable_error
 from ._utils.pref_vector import pref_vector_to_str_suffix, pref_vector_to_weighting
-from ._weighting_bases import Weighting
+from ._weighting_bases import GramianWeighting
 
 
-class DualProjWeighting(Weighting[PSDMatrix]):
+class DualProjWeighting(GramianWeighting):
     r"""
     :class:`~torchjd.aggregation._weighting_bases.Weighting` giving the weights of
     :class:`~torchjd.aggregation.DualProj`.
diff --git a/src/torchjd/aggregation/_gradvac.py b/src/torchjd/aggregation/_gradvac.py
@@ -10,10 +10,10 @@
 
 from ._aggregator_bases import GramianWeightedAggregator
 from ._utils.non_differentiable import raise_non_differentiable_error
-from ._weighting_bases import Weighting
+from ._weighting_bases import GramianWeighting
 
 
-class GradVacWeighting(Weighting[PSDMatrix], Stateful):
+class GradVacWeighting(GramianWeighting, Stateful):
     r"""
     :class:`~torchjd.aggregation._mixins.Stateful`
     :class:`~torchjd.aggregation._weighting_bases.Weighting` giving the weights of
diff --git a/src/torchjd/aggregation/_imtl_g.py b/src/torchjd/aggregation/_imtl_g.py
@@ -5,10 +5,10 @@
 
 from ._aggregator_bases import GramianWeightedAggregator
 from ._utils.non_differentiable import raise_non_differentiable_error
-from ._weighting_bases import Weighting
+from ._weighting_bases import GramianWeighting
 
 
-class IMTLGWeighting(Weighting[PSDMatrix]):
+class IMTLGWeighting(GramianWeighting):
     """
     :class:`~torchjd.aggregation._weighting_bases.Weighting` giving the weights of
     :class:`~torchjd.aggregation.IMTLG`.
diff --git a/src/torchjd/aggregation/_krum.py b/src/torchjd/aggregation/_krum.py
@@ -5,10 +5,10 @@
 from torchjd._linalg import PSDMatrix
 
 from ._aggregator_bases import GramianWeightedAggregator
-from ._weighting_bases import Weighting
+from ._weighting_bases import GramianWeighting
 
 
-class KrumWeighting(Weighting[PSDMatrix]):
+class KrumWeighting(GramianWeighting):
     """
     :class:`~torchjd.aggregation._weighting_bases.Weighting` giving the weights of
     :class:`~torchjd.aggregation.Krum`.
diff --git a/src/torchjd/aggregation/_mean.py b/src/torchjd/aggregation/_mean.py
@@ -1,13 +1,11 @@
 import torch
 from torch import Tensor
 
-from torchjd._linalg import Matrix
-
 from ._aggregator_bases import WeightedAggregator
-from ._weighting_bases import Weighting
+from ._weighting_bases import MatrixWeighting
 
 
-class MeanWeighting(Weighting[Matrix]):
+class MeanWeighting(MatrixWeighting):
     r"""
     :class:`~torchjd.aggregation._weighting_bases.Weighting` that gives the weights
     :math:`\begin{bmatrix} \frac{1}{m} & \dots & \frac{1}{m} \end{bmatrix}^T \in
diff --git a/src/torchjd/aggregation/_mgda.py b/src/torchjd/aggregation/_mgda.py
@@ -4,10 +4,10 @@
 from torchjd._linalg import PSDMatrix
 
 from ._aggregator_bases import GramianWeightedAggregator
-from ._weighting_bases import Weighting
+from ._weighting_bases import GramianWeighting
 
 
-class MGDAWeighting(Weighting[PSDMatrix]):
+class MGDAWeighting(GramianWeighting):
     r"""
     :class:`~torchjd.aggregation._weighting_bases.Weighting` giving the weights of
     :class:`~torchjd.aggregation.MGDA`.
diff --git a/src/torchjd/aggregation/_nash_mtl.py b/src/torchjd/aggregation/_nash_mtl.py
@@ -1,11 +1,10 @@
 # Partly adapted from https://github.com/AvivNavon/nash-mtl — MIT License, Copyright (c) 2022 Aviv Navon.
 # See NOTICES for the full license text.
 
-from torchjd._linalg import Matrix
 from torchjd.aggregation._mixins import Stateful
 
 from ._utils.check_dependencies import check_dependencies_are_installed
-from ._weighting_bases import Weighting
+from ._weighting_bases import MatrixWeighting
 
 check_dependencies_are_installed(["cvxpy", "ecos"])
 
@@ -21,7 +20,7 @@
 from ._utils.non_differentiable import raise_non_differentiable_error
 
 
-class _NashMTLWeighting(Weighting[Matrix], Stateful):
+class _NashMTLWeighting(MatrixWeighting, Stateful):
     """
     :class:`~torchjd.aggregation._mixins.Stateful` :class:`~torchjd.aggregation.Weighting` that
     extracts weights using the step decision of Algorithm 1 of `Multi-Task Learning as a Bargaining
diff --git a/src/torchjd/aggregation/_pcgrad.py b/src/torchjd/aggregation/_pcgrad.py
@@ -7,10 +7,10 @@
 
 from ._aggregator_bases import GramianWeightedAggregator
 from ._utils.non_differentiable import raise_non_differentiable_error
-from ._weighting_bases import Weighting
+from ._weighting_bases import GramianWeighting
 
 
-class PCGradWeighting(Weighting[PSDMatrix]):
+class PCGradWeighting(GramianWeighting):
     """
     :class:`~torchjd.aggregation._weighting_bases.Weighting` giving the weights of
     :class:`~torchjd.aggregation.PCGrad`.
diff --git a/src/torchjd/aggregation/_random.py b/src/torchjd/aggregation/_random.py
@@ -2,13 +2,11 @@
 from torch import Tensor
 from torch.nn import functional as F
 
-from torchjd._linalg import Matrix
-
 from ._aggregator_bases import WeightedAggregator
-from ._weighting_bases import Weighting
+from ._weighting_bases import MatrixWeighting
 
 
-class RandomWeighting(Weighting[Matrix]):
+class RandomWeighting(MatrixWeighting):
     """
     :class:`~torchjd.aggregation._weighting_bases.Weighting` that generates positive random weights
     at each call.
diff --git a/src/torchjd/aggregation/_sum.py b/src/torchjd/aggregation/_sum.py
@@ -1,13 +1,11 @@
 import torch
 from torch import Tensor
 
-from torchjd._linalg import Matrix
-
 from ._aggregator_bases import WeightedAggregator
-from ._weighting_bases import Weighting
+from ._weighting_bases import MatrixWeighting
 
 
-class SumWeighting(Weighting[Matrix]):
+class SumWeighting(MatrixWeighting):
     r"""
     :class:`~torchjd.aggregation._weighting_bases.Weighting` that gives the weights
     :math:`\begin{bmatrix} 1 & \dots & 1 \end{bmatrix}^T \in \mathbb{R}^m`.
diff --git a/src/torchjd/aggregation/_upgrad.py b/src/torchjd/aggregation/_upgrad.py
@@ -8,10 +8,10 @@
 from ._utils.dual_cone import SUPPORTED_SOLVER, project_weights
 from ._utils.non_differentiable import raise_non_differentiable_error
 from ._utils.pref_vector import pref_vector_to_str_suffix, pref_vector_to_weighting
-from ._weighting_bases import Weighting
+from ._weighting_bases import GramianWeighting
 
 
-class UPGradWeighting(Weighting[PSDMatrix]):
+class UPGradWeighting(GramianWeighting):
     r"""
     :class:`~torchjd.aggregation._weighting_bases.Weighting` giving the weights of
     :class:`~torchjd.aggregation.UPGrad`.
diff --git a/src/torchjd/aggregation/_weighting_bases.py b/src/torchjd/aggregation/_weighting_bases.py

Original file line number	Diff line number	Diff line change
`@@ -4,3 +4,4 @@ ConFIG`
`4`	`4`	`======`
`5`	`5`
`6`	`6`	`.. autoclass:: torchjd.aggregation.ConFIG`
	`7`	`+ :members: __call__`
Original file line number	Diff line number	Diff line change
`@@ -4,3 +4,4 @@ Flattening`
`4`	`4`	`==========`
`5`	`5`
`6`	`6`	`.. autoclass:: torchjd.aggregation.Flattening`
	`7`	`+ :members: __call__`
Original file line number	Diff line number	Diff line change
`@@ -4,3 +4,4 @@ GradDrop`
`4`	`4`	`========`
`5`	`5`
`6`	`6`	`.. autoclass:: torchjd.aggregation.GradDrop`
	`7`	`+ :members: __call__`