Add and use JacobianBasedGramianComputerWithoutCrossTerms

Author: ValerianRey

src/torchjd/autogram/_engine.py

@@ -8,7 +8,7 @@
 from ._gramian_accumulator import GramianAccumulator
 from ._gramian_computer import (
     GramianComputer,
-    JacobianBasedGramianComputerWithCrossTerms,
+    JacobianBasedGramianComputerWithoutCrossTerms,
     LinearBasedGramianComputer,
 )
 from ._gramian_utils import movedim_gramian, reshape_gramian
@@ -213,10 +213,10 @@ def _make_gramian_computer(self, module: nn.Module) -> GramianComputer:
                 gramian_computer = LinearBasedGramianComputer(module)
             else:
                 jacobian_computer = FunctionalJacobianComputer(module)
-                gramian_computer = JacobianBasedGramianComputerWithCrossTerms(jacobian_computer)
+                gramian_computer = JacobianBasedGramianComputerWithoutCrossTerms(jacobian_computer)
         else:
             jacobian_computer = AutogradJacobianComputer(module)
-            gramian_computer = JacobianBasedGramianComputerWithCrossTerms(jacobian_computer)
+            gramian_computer = JacobianBasedGramianComputerWithoutCrossTerms(jacobian_computer)
 
         return gramian_computer
 

src/torchjd/autogram/_gramian_computer.py

@@ -81,6 +81,25 @@ def __call__(
             return None
 
 
+class JacobianBasedGramianComputerWithoutCrossTerms(JacobianBasedGramianComputer):
+    """
+    Stateful JacobianBasedGramianComputer that directly returning the gramian without considering
+    cross-terms (except intra-module cross-terms).
+    """
+
+    def __call__(
+        self,
+        rg_outputs: tuple[Tensor, ...],
+        grad_outputs: tuple[Tensor, ...],
+        args: tuple[PyTree, ...],
+        kwargs: dict[str, PyTree],
+    ) -> Optional[Tensor]:
+        """Compute what we can for a module and optionally return the gramian if it's ready."""
+
+        jacobian_matrix = self.jacobian_computer(rg_outputs, grad_outputs, args, kwargs)
+        return self._to_gramian(jacobian_matrix)
+
+
 class ModuleBasedGramianComputer(GramianComputer, ABC):
     def __init__(self, module: nn.Module):
         self.module = module