refactor(linalg): Make Matrix and PSDMatrix public (#673)

- Introduces a new public `torchjd.linalg` package exposing `Matrix` and
`PSDMatrix` (the rest of `_linalg` stays protected)
- Makes `MatrixWeighting` and `GramianWeighting` protected. These
classes are still used to specify the docstring of the `__call__`
methods of the aggregators, but the user only sees those aggregators as
`Weighting[Matrix]` and `Weighting[PSDMatrix]`, respectively. The
`MatrixWeighting` and `GramianWeighting` classes really just bring
updated docstrings, that's all.
- Makes the public type of the gramian_weighting of
GramianWeightedAggregator be Weighting[PSDMatrix] instead of
GramianWeighting, so that #669 can work. Similar with weighting of
WeightedAggregator being Weighting[Matrix].
- Expands docstrings on `Matrix` and `PSDMatrix` with Jacobian and
Gramian examples; adds Sphinx documentation under a new **linalg**
section in the API Reference

Author: ValerianRey

CHANGELOG.md

@@ -10,9 +10,12 @@ 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.
-  They can be extended to easily implement custom `Weighting`s and `Aggregator`s.
+- Made `WeightedAggregator` and `GramianWeightedAggregator` public. These abstract base classes are
+  now importable from `torchjd.aggregation` and documented. They can be extended to easily implement
+  custom `Aggregator`s.
+- Made `Matrix` and `PSDMatrix` public. These type annotation classes are now importable from
+  `torchjd.linalg` and documented. Users can now subclass `Weighting[Matrix]` or
+  `Weighting[PSDMatrix]` to implement custom `Weighting`s.
 - 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

@@ -19,12 +19,6 @@ Abstract base classes
 .. autoclass:: torchjd.aggregation.Weighting
     :members: __call__
 
-.. autoclass:: torchjd.aggregation.MatrixWeighting
-    :members: __call__
-
-.. autoclass:: torchjd.aggregation.GramianWeighting
-    :members: __call__
-
 .. autoclass:: torchjd.aggregation.GeneralizedWeighting
     :members: __call__
 

docs/source/docs/linalg/index.rst

@@ -0,0 +1,12 @@
+linalg
+======
+
+.. automodule:: torchjd.linalg
+   :no-members:
+
+.. toctree::
+    :hidden:
+    :maxdepth: 1
+
+    matrix.rst
+    psd_matrix.rst

docs/source/docs/linalg/matrix.rst

@@ -0,0 +1,6 @@
+:hide-toc:
+
+Matrix
+======
+
+.. autoclass:: torchjd.linalg.Matrix

docs/source/docs/linalg/psd_matrix.rst

@@ -0,0 +1,6 @@
+:hide-toc:
+
+PSDMatrix
+=========
+
+.. autoclass:: torchjd.linalg.PSDMatrix

docs/source/index.rst

@@ -70,3 +70,4 @@ TorchJD is open-source, under MIT License. The source code is available on
     docs/autogram/index.rst
     docs/autojac/index.rst
     docs/aggregation/index.rst
+    docs/linalg/index.rst

src/torchjd/_linalg/_matrix.py

@@ -8,7 +8,16 @@
 
 
 class Matrix(Tensor):
-    """Tensor with exactly 2 dimensions."""
+    """
+    Tensor with exactly 2 dimensions.
+
+    Common examples include the Jacobian matrix J of shape ``[m, n]``, where m is the number of
+    objectives and n is the number of model parameters, and the Gramian of the Jacobian
+    G = J J^T of shape ``[m, m]``.
+
+    .. note::
+        This class should never be instantiated. It is only used for static type checking.
+    """
 
 
 class PSDTensor(Tensor):
@@ -20,7 +29,15 @@ class PSDTensor(Tensor):
 
 
 class PSDMatrix(PSDTensor, Matrix):
-    """Positive semi-definite matrix."""
+    """
+    Positive semi-definite matrix.
+
+    A common example is the Gramian of the Jacobian G = J J^T of shape ``[m, m]``, where J is a
+    Jacobian matrix of shape ``[m, n]``.
+
+    .. note::
+        This class should never be instantiated. It is only used for static type checking.
+    """
 
 
 def is_matrix(t: Tensor) -> TypeGuard[Matrix]:

src/torchjd/aggregation/__init__.py

@@ -9,8 +9,9 @@
 
 .. note::
     Most aggregators rely on computing the Gramian of the Jacobian, extracting a vector of weights
-    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
+    from this Gramian using a :class:`~torchjd.aggregation.Weighting`
+    [:class:`~torchjd.linalg.PSDMatrix`], 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`
@@ -80,7 +81,7 @@
 from ._utils.check_dependencies import (
     OptionalDepsNotInstalledError as _OptionalDepsNotInstalledError,
 )
-from ._weighting_bases import GeneralizedWeighting, GramianWeighting, MatrixWeighting, Weighting
+from ._weighting_bases import GeneralizedWeighting, Weighting
 
 __all__ = [
     "Aggregator",
@@ -97,12 +98,10 @@
     "GradVac",
     "GradVacWeighting",
     "GramianWeightedAggregator",
-    "GramianWeighting",
     "IMTLG",
     "IMTLGWeighting",
     "Krum",
     "KrumWeighting",
-    "MatrixWeighting",
     "Mean",
     "MeanWeighting",
     "MGDA",

src/torchjd/aggregation/_aggregator_bases.py

@@ -1,11 +1,11 @@
 from abc import ABC, abstractmethod
-from typing import cast
 
 from torch import Tensor, nn
 
-from torchjd._linalg import Matrix, compute_gramian, is_matrix
+from torchjd._linalg import compute_gramian, is_matrix
+from torchjd.linalg import Matrix, PSDMatrix
 
-from ._weighting_bases import GramianWeighting, MatrixWeighting
+from ._weighting_bases import Weighting
 
 
 class Aggregator(nn.Module, ABC):
@@ -48,12 +48,12 @@ def __str__(self) -> str:
 class WeightedAggregator(Aggregator):
     """
     Aggregator that combines the rows of the input Jacobian matrix with weights given by applying a
-    :class:`~torchjd.aggregation.MatrixWeighting` to it.
+    :class:`~torchjd.aggregation.Weighting` [:class:`~torchjd.linalg.Matrix`] to it.
 
     :param weighting: The object responsible for extracting the vector of weights from the matrix.
     """
 
-    def __init__(self, weighting: MatrixWeighting) -> None:
+    def __init__(self, weighting: Weighting[Matrix]) -> None:
         super().__init__()
         self.weighting = weighting
 
@@ -76,12 +76,13 @@ def forward(self, matrix: Matrix, /) -> Tensor:
 class GramianWeightedAggregator(WeightedAggregator):
     """
     :class:`~torchjd.aggregation.WeightedAggregator` that computes the gramian of the input
-    Jacobian matrix before applying a :class:`~torchjd.aggregation.GramianWeighting` to it.
+    Jacobian matrix before applying a :class:`~torchjd.aggregation.Weighting`
+    [:class:`~torchjd.linalg.PSDMatrix`] to it.
 
     :param gramian_weighting: The object responsible for extracting the vector of weights from the
         gramian.
     """
 
-    def __init__(self, gramian_weighting: GramianWeighting) -> None:
-        super().__init__(cast(MatrixWeighting, gramian_weighting << compute_gramian))
+    def __init__(self, gramian_weighting: Weighting[PSDMatrix]) -> None:
+        super().__init__(gramian_weighting << compute_gramian)
         self.gramian_weighting = gramian_weighting

src/torchjd/aggregation/_aligned_mtl.py

@@ -7,20 +7,20 @@
 import torch
 from torch import Tensor
 
-from torchjd._linalg import PSDMatrix
+from torchjd.linalg import PSDMatrix
 
 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 GramianWeighting
+from ._weighting_bases import _GramianWeighting
 
 SUPPORTED_SCALE_MODE: TypeAlias = Literal["min", "median", "rmse"]
 
 
-class AlignedMTLWeighting(GramianWeighting):
+class AlignedMTLWeighting(_GramianWeighting):
     r"""
-    :class:`~torchjd.aggregation.GramianWeighting` giving the weights of
-    :class:`~torchjd.aggregation.AlignedMTL`.
+    :class:`~torchjd.aggregation.Weighting` [:class:`~torchjd.linalg.PSDMatrix`]
+    giving the weights of :class:`~torchjd.aggregation.AlignedMTL`.
 
     :param pref_vector: The preference vector to use. If not provided, defaults to
         :math:`\begin{bmatrix} \frac{1}{m} & \dots & \frac{1}{m} \end{bmatrix}^T \in \mathbb{R}^m`.