fix(autogram): Raise error on forward pass without backward pass

CHANGELOG.md

@@ -18,6 +18,9 @@ changelog does not include internal changes that do not affect the user.
   matrix are almost equal. Rounding errors could make the squared distance between such rows
   slightly negative, giving a `nan` distance that was then ignored when computing the scores.
   Squared distances are now clamped to be non-negative before taking the square root.
+- Fixed `Engine.compute_gramian` silently ignoring the contribution of a module when one of its
+  forward passes had no associated backward pass (e.g. because its output was detached). It now
+  raises a `ValueError` instead.
 
 ## [0.17.1] - 2026-09-23
 

src/torchjd/autogram/_engine.py

@@ -326,6 +326,12 @@ def differentiation(_grad_output: Tensor) -> tuple[Tensor, ...]:
             grad_output = torch.ones_like(output)
             _ = differentiation(grad_output)
 
+        if any(computer.has_pending_jacobian() for computer in self._gramian_computers.values()):
+            raise ValueError(
+                "Some modules had a forward pass without an associated backward pass (e.g. because "
+                "its output was detached)."
+            )
+
         # If the gramian were None, then leaf_targets would be empty, so autograd.grad would
         # have failed. So gramian is necessarily a valid Tensor here.
         gramian = cast(PSDMatrix, self._gramian_accumulator.gramian)

src/torchjd/autogram/_gramian_computer.py

@@ -28,6 +28,12 @@ def track_forward_call(self) -> None:
     def reset(self) -> None:
         """Reset state if any. Necessary in some implementations."""
 
+    @abstractmethod
+    def has_pending_jacobian(self) -> bool:
+        """
+        Whether some Jacobian was computed but its Gramian was never returned. Necessary in some implementations.
+        """
+
 
 class JacobianBasedGramianComputer(GramianComputer, ABC):
     def __init__(self, jacobian_computer: JacobianComputer) -> None:
@@ -52,6 +58,9 @@ def reset(self) -> None:
     def track_forward_call(self) -> None:
         self.remaining_counter += 1
 
+    def has_pending_jacobian(self) -> bool:
+        return self.summed_jacobian is not None
+
     def __call__(
         self,
         rg_outputs: tuple[Tensor, ...],
@@ -72,6 +81,6 @@ def __call__(
 
         if self.remaining_counter == 0:
             gramian = compute_gramian(self.summed_jacobian)
-            del self.summed_jacobian
+            self.summed_jacobian = None
             return gramian
         return None

tests/unit/autogram/test_engine.py

@@ -570,3 +570,22 @@ def test_batched_non_batched_equivalence_2(factory: ModuleFactory, batch_size: i
     gramian_none = engine_none.compute_gramian(losses_none)
 
     assert_close(gramian_0, gramian_none, rtol=1e-4, atol=1e-5)
+
+
+@mark.parametrize("batch_dim", [0, None])
+def test_forward_pass_without_backward_pass_raises(batch_dim: int | None) -> None:
+    """
+    Tests that compute_gramian raises an error when a module has a forward pass without an
+    associated backward pass, instead of silently ignoring the contribution of this module.
+    """
+
+    encoder = ModuleFactory(Linear, 3, 3)()
+    head = ModuleFactory(Linear, 3, 1)()
+    engine = Engine(encoder, head, batch_dim=batch_dim)
+    input = randn_((4, 3))
+
+    target = encoder(input + 0.1).detach()  # No backward pass will go through this forward pass
+    losses = (head(encoder(input)) - target.sum(dim=1, keepdim=True)).squeeze(dim=1) ** 2
+
+    with pytest.raises(ValueError):
+        engine.compute_gramian(losses)