diff --git a/CHANGELOG.md b/CHANGELOG.md index 5b0fa6cca..170f05302 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/torchjd/autogram/_engine.py b/src/torchjd/autogram/_engine.py index 8a8633df9..766e191e9 100644 --- a/src/torchjd/autogram/_engine.py +++ b/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) diff --git a/src/torchjd/autogram/_gramian_computer.py b/src/torchjd/autogram/_gramian_computer.py index e44884971..938012738 100644 --- a/src/torchjd/autogram/_gramian_computer.py +++ b/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 diff --git a/tests/unit/autogram/test_engine.py b/tests/unit/autogram/test_engine.py index 84fb1f5f4..6cdde8cc7 100644 --- a/tests/unit/autogram/test_engine.py +++ b/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)