Skip to content

Fix TransformerBridge backward hook cleanup#1324

Open
SamuelePunzo wants to merge 2 commits into
TransformerLensOrg:mainfrom
SamuelePunzo:fix/1323-transformerbridge-bwd-hook-cleanup
Open

Fix TransformerBridge backward hook cleanup#1324
SamuelePunzo wants to merge 2 commits into
TransformerLensOrg:mainfrom
SamuelePunzo:fix/1323-transformerbridge-bwd-hook-cleanup

Conversation

@SamuelePunzo
Copy link
Copy Markdown

@SamuelePunzo SamuelePunzo commented May 25, 2026

Fix TransformerBridge backward-hook cleanup so temporary backward hooks do not leak after context exit.

HookPoint.remove_hooks() defaults to forward-only cleanup, and TransformerBridge was calling it without an explicit direction in several cleanup paths. That
meant non-permanent backward hooks could remain attached after with model.hooks(...) exited, and reset_hooks() on bridge components could also miss backward
hooks.

This change:

  • cleans up both directions in generalized component hook removal
  • cleans up the specific direction that was added in TransformerBridge.hooks(...)
  • cleans up the specific direction that was added in TransformerBridge.run_with_hooks(...)
  • preserves unrelated backward hooks during run_with_cache(...) teardown by cleaning up only the forward cache hooks it adds
  • adds regression tests for backward-hook cleanup on context exit, reset_hooks(), and run_with_cache(...)

This fixes the bug reported in the linked issue. No additional dependencies are required.

Fixes #1323

Type of change

  • Bug fix (non-breaking change which fixes an issue)

Screenshots

Not applicable.

Checklist:

  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug Report] TransformerBridge leaks non-permanent backward hooks after hooks(...) context exit

1 participant