refactor(aggregation)!: Remove generalized gramians (#692)

* `Engine.compute_gramian` now always returns a flat `[m, m]` gramian, regardless of the output shape
* Remove `GeneralizedWeighting`, and `Flattening` — they are no longer needed
* Update the IWMTL example to use `UPGradWeighting` directly and reshape the weights before calling `backward`
* Add a migration guide in `CHANGELOG.md`
---------
Co-authored-by: Pierre Quinton <pierre.quinton@epfl.ch>

Author: ValerianRey

CHANGELOG.md

@@ -20,6 +20,29 @@ changelog does not include internal changes that do not affect the user.
   (installed with `pip install torchjd`) much lighter, but it means that users of `UPGrad` and
   `DualProj` now have to install the new optional dependency group `quadprog_projector` explicitly
   (with e.g. `pip install "torchjd[quadprog_projector]"`).
+- **BREAKING**: Removed entirely the concept of generalized Gramians. The `Engine.compute_gramian`
+  method now always returns a square matrix of shape `[m, m]`, where `m` is the total number of
+  elements of the ``output`` tensor (treating all dimensions uniformly). Previously, an output of
+  shape `[m1, m2]` would return a 4D generalized Gramian of shape `[m1, m2, m2, m1]`; it now
+  returns a `[m1 * m2, m1 * m2]` matrix.
+  This also removes `GeneralizedWeighting` and `Flattening`.
+  To update, replace `Flattening(weighting)` with a standard `Weighting` and reshape the resulting
+  weight vector yourself:
+  ```python
+  # Before
+  from torchjd.aggregation import Flattening, UPGradWeighting
+  weighting = Flattening(UPGradWeighting())
+  gramian = engine.compute_gramian(losses)  # shape: [m1, m2, m2, m1]
+  weights = weighting(gramian)              # shape: [m1, m2]
+  losses.backward(weights)
+
+  # After
+  from torchjd.aggregation import UPGradWeighting
+  weighting = UPGradWeighting()
+  gramian = engine.compute_gramian(losses)           # shape: [m1 * m2, m1 * m2]
+  weights = weighting(gramian).reshape(losses.shape) # shape: [m1, m2]
+  losses.backward(weights)
+  ```
 
 ## [0.11.0] - 2026-05-18
 

docs/source/docs/aggregation/flattening.rst

@@ -19,9 +19,6 @@ Abstract base classes
 .. autoclass:: torchjd.aggregation.Weighting
     :members: __call__
 
-.. autoclass:: torchjd.aggregation.GeneralizedWeighting
-    :members: __call__
-
 .. autoclass:: torchjd.aggregation.Stateful
     :members: reset
 
@@ -38,7 +35,6 @@ Abstract base classes
     cr_mogm.rst
     dualproj.rst
     fairgrad.rst
-    flattening.rst
     graddrop.rst
     gradvac.rst
     imtl_g.rst

docs/source/docs/aggregation/index.rst

@@ -16,7 +16,7 @@ The following example shows how to do that.
     from torch.nn import Linear, MSELoss, ReLU, Sequential
     from torch.optim import SGD
 
-    from torchjd.aggregation import Flattening, UPGradWeighting
+    from torchjd.aggregation import UPGradWeighting
     from torchjd.autogram import Engine
 
     shared_module = Sequential(Linear(10, 5), ReLU(), Linear(5, 3), ReLU())
@@ -30,7 +30,7 @@ The following example shows how to do that.
 
     optimizer = SGD(params, lr=0.1)
     mse = MSELoss(reduction="none")
-    weighting = Flattening(UPGradWeighting())
+    weighting = UPGradWeighting()
     engine = Engine(shared_module, batch_dim=0)
 
     inputs = torch.randn(8, 16, 10)  # 8 batches of 16 random input vectors of length 10
@@ -46,20 +46,19 @@ The following example shows how to do that.
         losses = torch.stack([mse(out1, target1), mse(out2, target2)], dim=1)  # shape: [16, 2]
 
         # Compute the gramian (inner products between pairs of gradients of the losses)
-        gramian = engine.compute_gramian(losses)  # shape: [16, 2, 2, 16]
+        gramian = engine.compute_gramian(losses)  # shape: [32, 32]
 
         # Obtain the weights that lead to no conflict between reweighted gradients
-        weights = weighting(gramian)  # shape: [16, 2]
+        weights = weighting(gramian)  # shape: [32]
 
         # Do the standard backward pass, but weighted using the obtained weights
-        losses.backward(weights)
+        losses.backward(weights.reshape(losses.shape))
         optimizer.step()
         optimizer.zero_grad()
 
 .. note::
-    In this example, the tensor of losses is a matrix rather than a vector. The gramian is thus a
-    4D tensor rather than a matrix, and a
-    :class:`~torchjd.aggregation._weighting_bases.GeneralizedWeighting`, such as
-    :class:`~torchjd.aggregation._flattening.Flattening`, has to be used to extract a matrix of
-    weights from it. More information about ``GeneralizedWeighting`` can be found in the
-    :doc:`../../docs/aggregation/index` page.
+    In this example, the tensor of losses is a matrix of shape ``[16, 2]`` (16 samples, 2 tasks).
+    The autogram engine flattens this into a vector of ``m = 16 × 2 = 32`` objectives, so the
+    Gramian has shape ``[32, 32]``. A standard :class:`~torchjd.aggregation.Weighting` is then used
+    to extract a vector of 32 weights, which is reshaped back to ``[16, 2]`` before being passed to
+    :meth:`~torch.Tensor.backward`.

docs/source/examples/iwmtl.rst

@@ -36,28 +36,6 @@
 >>> weights = weighting(gramian)
 >>> weights
 tensor([1.1109, 0.7894])
-
-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.GeneralizedWeighting>` to extract
-a tensor of weights (of shape ``[m_1, ..., m_k]``) from such a generalized Gramian. The simplest
-: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.
-
->>> from torch import ones
->>> from torchjd.aggregation import Flattening, UPGradWeighting
->>>
->>> weighting = Flattening(UPGradWeighting())
->>> # Generate a generalized Gramian filled with ones, for the sake of the example
->>> generalized_gramian = ones((2, 3, 3, 2))
->>> weights = weighting(generalized_gramian)
->>> weights
-tensor([[0.1667, 0.1667, 0.1667],
-        [0.1667, 0.1667, 0.1667]])
 """
 
 from ._aggregator_bases import Aggregator, GramianWeightedAggregator, WeightedAggregator
@@ -68,7 +46,6 @@
 from ._cr_mogm import CRMOGMWeighting
 from ._dualproj import DualProj, DualProjWeighting
 from ._fairgrad import FairGrad, FairGradWeighting
-from ._flattening import Flattening
 from ._graddrop import GradDrop
 from ._gradvac import GradVac, GradVacWeighting
 from ._imtl_g import IMTLG, IMTLGWeighting
@@ -82,7 +59,7 @@
 from ._sum import Sum, SumWeighting
 from ._trimmed_mean import TrimmedMean
 from ._upgrad import UPGrad, UPGradWeighting
-from ._weighting_bases import GeneralizedWeighting, Weighting
+from ._weighting_bases import Weighting
 
 __all__ = [
     "Aggregator",
@@ -98,8 +75,6 @@
     "DualProjWeighting",
     "FairGrad",
     "FairGradWeighting",
-    "Flattening",
-    "GeneralizedWeighting",
     "GradDrop",
     "GradVac",
     "GradVacWeighting",

src/torchjd/aggregation/__init__.py

@@ -6,7 +6,6 @@
 
 from torch import Tensor, nn
 
-from torchjd._linalg import PSDTensor, is_psd_tensor
 from torchjd.linalg import Matrix, PSDMatrix
 
 _T = TypeVar("_T", contravariant=True, bound=Tensor)
@@ -76,30 +75,3 @@ def __call__(self, gramian: Tensor, /) -> Tensor:
         :param gramian: The Gramian from which the weights must be extracted.
         """
         return super().__call__(gramian)
-
-
-class GeneralizedWeighting(nn.Module, ABC):
-    r"""
-    Abstract base class for all weightings that operate on generalized Gramians. It has the role of
-    extracting a tensor of weights of dimension :math:`m_1 \times \dots \times m_k` from a
-    generalized Gramian of dimension
-    :math:`m_1 \times \dots \times m_k \times m_k \times \dots \times m_1`.
-    """
-
-    def __init__(self) -> None:
-        super().__init__()
-
-    @abstractmethod
-    def forward(self, generalized_gramian: PSDTensor, /) -> Tensor:
-        """Computes the vector of weights from the input generalized Gramian."""
-
-    def __call__(self, generalized_gramian: Tensor, /) -> Tensor:
-        """
-        Computes the tensor of weights from the input generalized Gramian and applies all registered
-        hooks.
-
-        :param generalized_gramian: The tensor from which the weights must be extracted.
-        """
-
-        assert is_psd_tensor(generalized_gramian)
-        return super().__call__(generalized_gramian)

src/torchjd/aggregation/_flattening.py

@@ -4,7 +4,7 @@
 from torch import Tensor, nn, vmap
 from torch.autograd.graph import get_gradient_edge
 
-from torchjd._linalg import movedim, reshape
+from torchjd._linalg import flatten, movedim, reshape
 from torchjd.linalg import PSDMatrix
 
 from ._edge_registry import EdgeRegistry
@@ -246,28 +246,18 @@ def compute_gramian(self, output: Tensor, /) -> Tensor:
         Computes the Gramian of the Jacobian of ``output`` with respect to the direct parameters of
         all ``modules``.
 
-        :param output: The tensor of arbitrary shape to differentiate. The shape of the returned
-            Gramian depends on the shape of this output.
-
-        .. note::
-            This function doesn't require ``output`` to be a vector. For example, if ``output`` is
-            a matrix of shape :math:`[m_1, m_2]`, its Jacobian :math:`J` with respect to the
-            parameters will be of shape :math:`[m_1, m_2, n]`, where :math:`n` is the number of
-            parameters in the model. This is what we call a `generalized Jacobian`. The
-            corresponding Gramian :math:`G = J J^\top` will be of shape
-            :math:`[m_1, m_2, m_2, m_1]`. This is what we call a `generalized Gramian`. The number
-            of dimensions of the returned generalized Gramian will always be twice that of the
-            ``output``.
+        :param output: The tensor to differentiate. Its elements are treated as a flat vector of
+            :math:`m` objectives (where :math:`m` is the total number of elements of ``output``),
+            so the returned Gramian always has shape :math:`[m, m]`.
 
             A few examples:
-                - 0D (scalar) ``output``: 0D Gramian (this can be used to efficiently compute the
-                  squared norm of the gradient of ``output``).
-                - 1D (vector) ``output``: 2D Gramian (this is the standard setting of Jacobian
-                  descent).
-                - 2D (matrix) ``output``: 4D Gramian (this can be used for :doc:`Instance-Wise
-                  Multi-Task Learning (IWMTL) <../../examples/iwmtl>`, as each sample in the batch
-                  has one loss per task).
-                - etc.
+                - Scalar ``output``: :math:`1\times 1` Gramian (this can be used to efficiently
+                  compute the squared norm of the gradient of ``output``).
+                - Vector ``output`` of dimension :math:`m`: :math:`m \times m` Gramian (this is the
+                  standard setting of Jacobian descent).
+                - Matrix ``output`` of dimension :math:`m_1\times m_2`: :math:`m_1 m_2 \times m_1 m_2`
+                  Gramian (this can be used for :doc:`Instance-Wise Multi-Task Learning (IWMTL)
+                  <../../examples/iwmtl>`, as each sample in the batch has one loss per task).
         """
 
         if self._batch_dim is not None:
@@ -305,12 +295,11 @@ def compute_gramian(self, output: Tensor, /) -> Tensor:
             for gramian_computer in self._gramian_computers.values():
                 gramian_computer.reset()
 
-        unordered_gramian = reshape(square_gramian, ordered_shape)
-
         if self._batch_dim is not None:
-            gramian = movedim(unordered_gramian, [-1], [self._batch_dim])
+            unordered_gramian = reshape(square_gramian, ordered_shape)
+            gramian = flatten(movedim(unordered_gramian, [-1], [self._batch_dim]))
         else:
-            gramian = unordered_gramian
+            gramian = square_gramian
 
         return gramian