From 355aa4952a90495cd1184d3b5ce76b3339e87ab7 Mon Sep 17 00:00:00 2001 From: ppraneth Date: Mon, 5 Oct 2026 18:42:55 +0530 Subject: [PATCH 1/2] minor fix --- CHANGELOG.md | 3 +++ src/torchjd/autogram/_engine.py | 5 +++++ src/torchjd/autogram/_gramian_computer.py | 11 ++++++++++- tests/unit/autogram/test_engine.py | 19 +++++++++++++++++++ 4 files changed, 37 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5b0fa6cca..4cce90ee3 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 this module + was also called in a forward pass that was not used to compute the output (e.g. a forward pass + whose output is 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..946e3f186 100644 --- a/src/torchjd/autogram/_engine.py +++ b/src/torchjd/autogram/_engine.py @@ -326,6 +326,11 @@ 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 were called in forward passes that were not used to compute `output`. " + ) + # 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..a56aa43a7 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_not_used_in_output_raises(batch_dim: int | None) -> None: + """ + Tests that compute_gramian raises an error when a module was called in a forward pass that is + not used to compute the output, 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() # Forward pass of encoder not used to compute losses + losses = (head(encoder(input)) - target.sum(dim=1, keepdim=True)).squeeze(dim=1) ** 2 + + with pytest.raises(ValueError): + engine.compute_gramian(losses) From eef3f15b4b0c2ea5e3f25604dc5b5052789ab5e1 Mon Sep 17 00:00:00 2001 From: ppraneth Date: Tue, 6 Oct 2026 08:21:12 +0530 Subject: [PATCH 2/2] Reword forward pass without backward pass --- CHANGELOG.md | 6 +++--- src/torchjd/autogram/_engine.py | 3 ++- tests/unit/autogram/test_engine.py | 8 ++++---- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4cce90ee3..170f05302 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,9 +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 this module - was also called in a forward pass that was not used to compute the output (e.g. a forward pass - whose output is detached). It now raises a `ValueError` instead. +- 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 946e3f186..766e191e9 100644 --- a/src/torchjd/autogram/_engine.py +++ b/src/torchjd/autogram/_engine.py @@ -328,7 +328,8 @@ def differentiation(_grad_output: Tensor) -> tuple[Tensor, ...]: if any(computer.has_pending_jacobian() for computer in self._gramian_computers.values()): raise ValueError( - "Some modules were called in forward passes that were not used to compute `output`. " + "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 diff --git a/tests/unit/autogram/test_engine.py b/tests/unit/autogram/test_engine.py index a56aa43a7..6cdde8cc7 100644 --- a/tests/unit/autogram/test_engine.py +++ b/tests/unit/autogram/test_engine.py @@ -573,10 +573,10 @@ def test_batched_non_batched_equivalence_2(factory: ModuleFactory, batch_size: i @mark.parametrize("batch_dim", [0, None]) -def test_forward_pass_not_used_in_output_raises(batch_dim: int | None) -> None: +def test_forward_pass_without_backward_pass_raises(batch_dim: int | None) -> None: """ - Tests that compute_gramian raises an error when a module was called in a forward pass that is - not used to compute the output, instead of silently ignoring the contribution of this module. + 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)() @@ -584,7 +584,7 @@ def test_forward_pass_not_used_in_output_raises(batch_dim: int | None) -> None: engine = Engine(encoder, head, batch_dim=batch_dim) input = randn_((4, 3)) - target = encoder(input + 0.1).detach() # Forward pass of encoder not used to compute losses + 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):