Add and use JacobianBasedGramianComputerWithoutCrossTerms

Author: ValerianRey

src/torchjd/autogram/_engine.py

@@ -6,7 +6,7 @@
 
 from ._edge_registry import EdgeRegistry
 from ._gramian_accumulator import GramianAccumulator
-from ._gramian_computer import GramianComputer, JacobianBasedGramianComputerWithCrossTerms
+from ._gramian_computer import GramianComputer, JacobianBasedGramianComputerWithoutCrossTerms
 from ._gramian_utils import movedim_gramian, reshape_gramian
 from ._jacobian_computer import (
     AutogradJacobianComputer,
@@ -205,9 +205,10 @@ def _make_gramian_computer(self, module: nn.Module) -> GramianComputer:
         jacobian_computer: JacobianComputer
         if self._batch_dim is not None:
             jacobian_computer = FunctionalJacobianComputer(module)
+            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

@@ -4,8 +4,6 @@
 from torch import Tensor
 from torch.utils._pytree import PyTree
 
-from torchjd.autogram._jacobian_computer import JacobianComputer
-
 
 class GramianComputer(ABC):
     @abstractmethod
@@ -34,24 +32,12 @@ def _to_gramian(jacobian: Tensor) -> Tensor:
         return jacobian @ jacobian.T
 
 
-class JacobianBasedGramianComputerWithCrossTerms(JacobianBasedGramianComputer):
+class JacobianBasedGramianComputerWithoutCrossTerms(JacobianBasedGramianComputer):
     """
-    Stateful JacobianBasedGramianComputer that waits for all usages to be counted before returning
-    the gramian.
+    Stateful JacobianBasedGramianComputer that directly returning the gramian without considering
+    cross-terms (except intra-module cross-terms).
     """
 
-    def __init__(self, jacobian_computer: JacobianComputer):
-        super().__init__(jacobian_computer)
-        self.remaining_counter = 0
-        self.summed_jacobian: Optional[Tensor] = None
-
-    def reset(self) -> None:
-        self.remaining_counter = 0
-        self.summed_jacobian = None
-
-    def track_forward_call(self) -> None:
-        self.remaining_counter += 1
-
     def __call__(
         self,
         rg_outputs: tuple[Tensor, ...],
@@ -62,17 +48,4 @@ def __call__(
         """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)
-
-        if self.summed_jacobian is None:
-            self.summed_jacobian = jacobian_matrix
-        else:
-            self.summed_jacobian += jacobian_matrix
-
-        self.remaining_counter -= 1
-
-        if self.remaining_counter == 0:
-            gramian = self._to_gramian(self.summed_jacobian)
-            del self.summed_jacobian
-            return gramian
-        else:
-            return None
+        return self._to_gramian(jacobian_matrix)