feat(aggregation): Make Aggregator and Weighting subclasses public (#670)

* Export `WeightedAggregator`, `GramianWeightedAggregator`,
`MatrixWeighting`, and `GramianWeighting` from `torchjd.aggregation`
(added to `__all__` and documented in `index.rst`)
* Update all aggregator and weighting docstrings to reference the most
specific public base class (e.g. `GramianWeightedAggregator` instead of
`Aggregator`, `GramianWeighting` instead of `Weighting`)
* Update internal imports that used the private `_aggregator_bases` path
to use the public `torchjd.aggregation` path
* Make the combine method protected
* Add changelog entry

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: ValerianRey

CHANGELOG.md

@@ -10,6 +10,8 @@ changelog does not include internal changes that do not affect the user.
 
 ### Added
 
+- Made `WeightedAggregator`, `GramianWeightedAggregator`, `MatrixWeighting`, and `GramianWeighting`
+  public. These abstract base classes are now importable from `torchjd.aggregation` and documented.
 - Added getters and setters for the constructor parameters of all aggregators and weightings, so
   that they can be changed after initialization. This includes: `pref_vector`,
   `norm_eps` and `reg_eps` in `UPGrad`, `UPGradWeighting`, `DualProj` and `DualProjWeighting`;

docs/source/docs/aggregation/index.rst

@@ -10,9 +10,21 @@ Abstract base classes
 .. autoclass:: torchjd.aggregation.Aggregator
     :members: __call__
 
+.. autoclass:: torchjd.aggregation.WeightedAggregator
+    :members: __call__
+
+.. autoclass:: torchjd.aggregation.GramianWeightedAggregator
+    :members: __call__
+
 .. autoclass:: torchjd.aggregation.Weighting
     :members: __call__
 
+.. autoclass:: torchjd.aggregation.MatrixWeighting
+    :members: __call__
+
+.. autoclass:: torchjd.aggregation.GramianWeighting
+    :members: __call__
+
 .. autoclass:: torchjd.aggregation.GeneralizedWeighting
     :members: __call__
 

src/torchjd/aggregation/__init__.py

@@ -1,23 +1,23 @@
 r"""
 When doing Jacobian descent, the Jacobian matrix has to be aggregated into a vector to store in the
 ``.grad`` fields of the model parameters. The
-:class:`~torchjd.aggregation._aggregator_bases.Aggregator` is responsible for these aggregations.
+:class:`~torchjd.aggregation.Aggregator` is responsible for these aggregations.
 
 When using the :doc:`autogram <../autogram/index>` engine, we rather need to extract a vector
 of weights from the Gramian of the Jacobian. The
-:class:`~torchjd.aggregation._weighting_bases.Weighting` is responsible for this.
+:class:`~torchjd.aggregation.Weighting` is responsible for this.
 
 .. note::
     Most aggregators rely on computing the Gramian of the Jacobian, extracting a vector of weights
-    from this Gramian using a :class:`~torchjd.aggregation._weighting_bases.Weighting`, and then
-    combining the rows of the Jacobian using these weights. For all of them, we provide both the
-    :class:`~torchjd.aggregation._aggregator_bases.Aggregator` interface (to be used in autojac) and
-    the :class:`~torchjd.aggregation._weighting_bases.Weighting` interface (to be used in autogram).
-    For the rest, we only provide the :class:`~torchjd.aggregation._aggregator_bases.Aggregator`
+    from this Gramian using a :class:`~torchjd.aggregation.GramianWeighting`, and then combining the
+    rows of the Jacobian using these weights. For all of them, we provide both the
+    :class:`~torchjd.aggregation.Aggregator` interface (to be used in autojac) and the
+    :class:`~torchjd.aggregation.Weighting` interface (to be used in autogram).
+    For the rest, we only provide the :class:`~torchjd.aggregation.Aggregator`
     interface -- they are not compatible with autogram.
 
-:class:`Aggregators <torchjd.aggregation._aggregator_bases.Aggregator>` and :class:`Weightings
-<torchjd.aggregation._weighting_bases.Weighting>` are callables that take a Jacobian matrix or a
+:class:`Aggregators <torchjd.aggregation.Aggregator>` and
+:class:`Weightings <torchjd.aggregation.Weighting>` are callables that take a Jacobian matrix or a
 Gramian matrix as inputs, respectively. The following example shows how to use UPGrad to either
 aggregate a Jacobian (of shape ``[m, n]``, where ``m`` is the number of objectives and ``n`` is the
 number of parameters), or obtain the weights from the Gramian of the Jacobian (of shape ``[m, m]``).
@@ -39,10 +39,10 @@
 When dealing with a more general tensor of objectives, of shape ``[m_1, ..., m_k]`` (i.e. not
 necessarily a simple vector), the Jacobian will be of shape ``[m_1, ..., m_k, n]``, and its Gramian
 will be called a `generalized Gramian`, of shape ``[m_1, ..., m_k, m_k, ..., m_1]``. One can use a
-:class:`GeneralizedWeighting<torchjd.aggregation._weighting_bases.GeneralizedWeighting>` to extract
+:class:`GeneralizedWeighting<torchjd.aggregation.GeneralizedWeighting>` to extract
 a tensor of weights (of shape ``[m_1, ..., m_k]``) from such a generalized Gramian. The simplest
-:class:`GeneralizedWeighting<torchjd.aggregation._weighting_bases.GeneralizedWeighting>` is
-:class:`Flattening<torchjd.aggregation._flattening.Flattening>`: it simply "flattens" the
+:class:`GeneralizedWeighting<torchjd.aggregation.GeneralizedWeighting>` is
+:class:`Flattening<torchjd.aggregation.Flattening>`: it simply "flattens" the
 generalized Gramian into a square Gramian matrix (of shape ``[m_1 * ... * m_k, m_1 * ... * m_k]``),
 applies a normal weighting to it to obtain a vector of weights, and returns the reshaped tensor of
 weights.
@@ -59,7 +59,7 @@
         [0.1667, 0.1667, 0.1667]])
 """
 
-from ._aggregator_bases import Aggregator
+from ._aggregator_bases import Aggregator, GramianWeightedAggregator, WeightedAggregator
 from ._aligned_mtl import AlignedMTL, AlignedMTLWeighting
 from ._config import ConFIG
 from ._constant import Constant, ConstantWeighting
@@ -80,7 +80,7 @@
 from ._utils.check_dependencies import (
     OptionalDepsNotInstalledError as _OptionalDepsNotInstalledError,
 )
-from ._weighting_bases import GeneralizedWeighting, Weighting
+from ._weighting_bases import GeneralizedWeighting, GramianWeighting, MatrixWeighting, Weighting
 
 __all__ = [
     "Aggregator",
@@ -96,10 +96,13 @@
     "GradDrop",
     "GradVac",
     "GradVacWeighting",
+    "GramianWeightedAggregator",
+    "GramianWeighting",
     "IMTLG",
     "IMTLGWeighting",
     "Krum",
     "KrumWeighting",
+    "MatrixWeighting",
     "Mean",
     "MeanWeighting",
     "MGDA",
@@ -114,6 +117,7 @@
     "TrimmedMean",
     "UPGrad",
     "UPGradWeighting",
+    "WeightedAggregator",
     "Weighting",
 ]
 

src/torchjd/aggregation/_aggregator_bases.py

@@ -1,10 +1,11 @@
 from abc import ABC, abstractmethod
+from typing import cast
 
 from torch import Tensor, nn
 
-from torchjd._linalg import Matrix, PSDMatrix, compute_gramian, is_matrix
+from torchjd._linalg import Matrix, compute_gramian, is_matrix
 
-from ._weighting_bases import Weighting
+from ._weighting_bases import GramianWeighting, MatrixWeighting
 
 
 class Aggregator(nn.Module, ABC):
@@ -46,18 +47,18 @@ def __str__(self) -> str:
 
 class WeightedAggregator(Aggregator):
     """
-    Aggregator that combines the rows of the input jacobian matrix with weights given by applying a
-    Weighting to it.
+    Aggregator that combines the rows of the input Jacobian matrix with weights given by applying a
+    :class:`~torchjd.aggregation.MatrixWeighting` to it.
 
     :param weighting: The object responsible for extracting the vector of weights from the matrix.
     """
 
-    def __init__(self, weighting: Weighting[Matrix]) -> None:
+    def __init__(self, weighting: MatrixWeighting) -> None:
         super().__init__()
         self.weighting = weighting
 
     @staticmethod
-    def combine(matrix: Matrix, weights: Tensor) -> Tensor:
+    def _combine(matrix: Matrix, weights: Tensor) -> Tensor:
         """
         Aggregates a matrix by making a linear combination of its rows, using the provided vector of
         weights.
@@ -68,19 +69,19 @@ def combine(matrix: Matrix, weights: Tensor) -> Tensor:
 
     def forward(self, matrix: Matrix, /) -> Tensor:
         weights = self.weighting(matrix)
-        vector = self.combine(matrix, weights)
+        vector = self._combine(matrix, weights)
         return vector
 
 
 class GramianWeightedAggregator(WeightedAggregator):
     """
-    WeightedAggregator that computes the gramian of the input jacobian matrix before applying a
-    Weighting to it.
+    :class:`~torchjd.aggregation.WeightedAggregator` that computes the gramian of the input
+    Jacobian matrix before applying a :class:`~torchjd.aggregation.GramianWeighting` to it.
 
     :param gramian_weighting: The object responsible for extracting the vector of weights from the
         gramian.
     """
 
-    def __init__(self, gramian_weighting: Weighting[PSDMatrix]) -> None:
-        super().__init__(gramian_weighting << compute_gramian)
+    def __init__(self, gramian_weighting: GramianWeighting) -> None:
+        super().__init__(cast(MatrixWeighting, gramian_weighting << compute_gramian))
         self.gramian_weighting = gramian_weighting

src/torchjd/aggregation/_aligned_mtl.py

@@ -19,7 +19,7 @@
 
 class AlignedMTLWeighting(GramianWeighting):
     r"""
-    :class:`~torchjd.aggregation._weighting_bases.Weighting` giving the weights of
+    :class:`~torchjd.aggregation.GramianWeighting` giving the weights of
     :class:`~torchjd.aggregation.AlignedMTL`.
 
     :param pref_vector: The preference vector to use. If not provided, defaults to
@@ -89,7 +89,7 @@ def _compute_balance_transformation(
 
 class AlignedMTL(GramianWeightedAggregator):
     r"""
-    :class:`~torchjd.aggregation._aggregator_bases.Aggregator` as defined in Algorithm 1 of
+    :class:`~torchjd.aggregation.GramianWeightedAggregator` as defined in Algorithm 1 of
     `Independent Component Alignment for Multi-Task Learning
     <https://openaccess.thecvf.com/content/CVPR2023/papers/Senushkin_Independent_Component_Alignment_for_Multi-Task_Learning_CVPR_2023_paper.pdf>`_.
 

src/torchjd/aggregation/_cagrad.py

@@ -20,7 +20,7 @@
 
 class CAGradWeighting(GramianWeighting):
     """
-    :class:`~torchjd.aggregation._weighting_bases.Weighting` giving the weights of
+    :class:`~torchjd.aggregation.GramianWeighting` giving the weights of
     :class:`~torchjd.aggregation.CAGrad`.
 
     :param c: The scale of the radius of the ball constraint.
@@ -94,7 +94,7 @@ def norm_eps(self, value: float) -> None:
 
 class CAGrad(GramianWeightedAggregator):
     """
-    :class:`~torchjd.aggregation._aggregator_bases.Aggregator` as defined in Algorithm 1 of
+    :class:`~torchjd.aggregation.GramianWeightedAggregator` as defined in Algorithm 1 of
     `Conflict-Averse Gradient Descent for Multi-task Learning
     <https://arxiv.org/pdf/2110.14048.pdf>`_.
 

src/torchjd/aggregation/_config.py

@@ -15,7 +15,7 @@
 
 class ConFIG(Aggregator):
     """
-    :class:`~torchjd.aggregation._aggregator_bases.Aggregator` as defined in Equation 2 of `ConFIG:
+    :class:`~torchjd.aggregation.Aggregator` as defined in Equation 2 of `ConFIG:
     Towards Conflict-free Training of Physics Informed Neural Networks
     <https://arxiv.org/pdf/2408.11104>`_.
 

src/torchjd/aggregation/_constant.py

@@ -7,7 +7,7 @@
 
 class ConstantWeighting(MatrixWeighting):
     """
-    :class:`~torchjd.aggregation._weighting_bases.Weighting` that returns constant, pre-determined
+    :class:`~torchjd.aggregation.MatrixWeighting` that returns constant, pre-determined
     weights.
 
     :param weights: The weights to return at each call.
@@ -37,7 +37,7 @@ def _check_matrix_shape(self, matrix: Tensor) -> None:
 
 class Constant(WeightedAggregator):
     """
-    :class:`~torchjd.aggregation._aggregator_bases.Aggregator` that makes a linear combination of
+    :class:`~torchjd.aggregation.WeightedAggregator` that makes a linear combination of
     the rows of the provided matrix, with constant, pre-determined weights.
 
     :param weights: The weights associated to the rows of the input matrices.

src/torchjd/aggregation/_dualproj.py

@@ -12,7 +12,7 @@
 
 class DualProjWeighting(GramianWeighting):
     r"""
-    :class:`~torchjd.aggregation._weighting_bases.Weighting` giving the weights of
+    :class:`~torchjd.aggregation.GramianWeighting` giving the weights of
     :class:`~torchjd.aggregation.DualProj`.
 
     :param pref_vector: The preference vector to use. If not provided, defaults to
@@ -78,7 +78,7 @@ def reg_eps(self, value: float) -> None:
 
 class DualProj(GramianWeightedAggregator):
     r"""
-    :class:`~torchjd.aggregation._aggregator_bases.Aggregator` that averages the rows of the input
+    :class:`~torchjd.aggregation.GramianWeightedAggregator` that averages the rows of the input
     matrix, and projects the result onto the dual cone of the rows of the matrix. This corresponds
     to the solution to Equation 11 of `Gradient Episodic Memory for Continual Learning
     <https://proceedings.neurips.cc/paper/2017/file/f87522788a2be2d171666752f97ddebb-Paper.pdf>`_.

src/torchjd/aggregation/_graddrop.py

@@ -15,7 +15,7 @@ def _identity(P: Tensor) -> Tensor:
 
 class GradDrop(Aggregator):
     """
-    :class:`~torchjd.aggregation._aggregator_bases.Aggregator` that applies the gradient combination
+    :class:`~torchjd.aggregation.Aggregator` that applies the gradient combination
     steps from GradDrop, as defined in lines 10 to 15 of Algorithm 1 of `Just Pick a Sign:
     Optimizing Deep Multitask Models with Gradient Sign Dropout
     <https://arxiv.org/pdf/2010.06808.pdf>`_.