factorise computing gramians

Author: PierreQuinton

src/torchjd/_utils/__init__.py

@@ -0,0 +1,3 @@
+from .compute_gramian import compute_gramian
+
+__all__ = ["compute_gramian"]

src/torchjd/_utils/compute_gramian.py

@@ -0,0 +1,15 @@
+import torch
+from torch import Tensor
+
+
+def compute_gramian(generalized_matrix: Tensor) -> Tensor:
+    """
+    Computes the `Gramian matrix <https://en.wikipedia.org/wiki/Gram_matrix>`_ of a given
+    generalized matrix. Specifically, this is equivalent to
+
+    matrix = generalized_matrix.reshape([generalized_matrix.shape[0], -1])
+    return matrix @ matrix.T
+    """
+    dims = list(range(1, generalized_matrix.ndim))
+    gramian = torch.tensordot(generalized_matrix, generalized_matrix, dims=(dims, dims))
+    return gramian

src/torchjd/aggregation/_aggregator_bases.py

@@ -2,7 +2,8 @@
 
 from torch import Tensor, nn
 
-from ._utils.gramian import compute_gramian
+from torchjd._utils import compute_gramian
+
 from ._weighting_bases import Matrix, PSDMatrix, Weighting
 
 

src/torchjd/aggregation/_utils/gramian.py

@@ -2,14 +2,6 @@
 from torch import Tensor
 
 
-def compute_gramian(matrix: Tensor) -> Tensor:
-    """
-    Computes the `Gramian matrix <https://en.wikipedia.org/wiki/Gram_matrix>`_ of a given matrix.
-    """
-
-    return matrix @ matrix.T
-
-
 def normalize(gramian: Tensor, eps: float) -> Tensor:
     """
     Normalizes the gramian `G=AA^T` with respect to the Frobenius norm of `A`.

src/torchjd/autogram/_gramian_computer.py

@@ -4,6 +4,7 @@
 from torch import Tensor
 from torch.utils._pytree import PyTree
 
+from torchjd._utils import compute_gramian
 from torchjd.autogram._jacobian_computer import JacobianComputer
 
 
@@ -29,10 +30,6 @@ class JacobianBasedGramianComputer(GramianComputer, ABC):
     def __init__(self, jacobian_computer):
         self.jacobian_computer = jacobian_computer
 
-    @staticmethod
-    def _to_gramian(jacobian: Tensor) -> Tensor:
-        return jacobian @ jacobian.T
-
 
 class JacobianBasedGramianComputerWithCrossTerms(JacobianBasedGramianComputer):
     """
@@ -71,7 +68,7 @@ def __call__(
         self.remaining_counter -= 1
 
         if self.remaining_counter == 0:
-            gramian = self._to_gramian(self.summed_jacobian)
+            gramian = compute_gramian(self.summed_jacobian)
             del self.summed_jacobian
             return gramian
         else:

tests/unit/aggregation/test_mgda.py

@@ -3,9 +3,9 @@
 from torch.testing import assert_close
 from utils.tensors import ones_, randn_
 
+from torchjd._utils.compute_gramian import compute_gramian
 from torchjd.aggregation import MGDA
 from torchjd.aggregation._mgda import MGDAWeighting
-from torchjd.aggregation._utils.gramian import compute_gramian
 
 from ._asserts import (
     assert_expected_structure,

tests/unit/aggregation/test_pcgrad.py

@@ -3,10 +3,10 @@
 from torch.testing import assert_close
 from utils.tensors import ones_, randn_
 
+from torchjd._utils.compute_gramian import compute_gramian
 from torchjd.aggregation import PCGrad
 from torchjd.aggregation._pcgrad import PCGradWeighting
 from torchjd.aggregation._upgrad import UPGradWeighting
-from torchjd.aggregation._utils.gramian import compute_gramian
 
 from ._asserts import assert_expected_structure, assert_non_differentiable
 from ._inputs import scaled_matrices, typical_matrices