Skip to content

fix(autogram): Remove reference cycles#424

Merged
PierreQuinton merged 5 commits intomainfrom
fix-autogram-ref-cycles
Sep 19, 2025
Merged

fix(autogram): Remove reference cycles#424
PierreQuinton merged 5 commits intomainfrom
fix-autogram-ref-cycles

Conversation

@ValerianRey
Copy link
Copy Markdown
Contributor

@ValerianRey ValerianRey commented Sep 18, 2025

  • Refactor ModuleHookManager
  • Refactor autograd Functions
  • Add finalizer in ModuleHookManager to unhook
  • Remove manual garbage collection

There will be some conflicts with #411 but I think we still want to merge this, and basically re-code this from scratch in #411 (it's quite easy now that we know what to do, maybe 1h work).

* Use a dedicated Hook class for readability and improved debugging
* Use BoolRef so that we can have a pointer to the boolean value in the Hook without having a pointer to the ModuleHookManager in the Hook (which is a reference cycle).
* Use the module given as input to Hook.__call__ rather than the module provided to hook_module. This avoids having to store a reference to the module in its hook (which is a reference cycle).
* Instead of using non-local variables, make them take these variables as input to forward
* This seems more standard practice, and fixes a reference cycle issue
* This solves the last reference cycle issue
* Remove garbage_collect marker
* Remove garbage_collect_if_marked fixture

This is not needed anymore since the garbage collector's job is much easier now
@ValerianRey ValerianRey added cc: fix Conventional commit type for bug fixes of the actual library (changes to src). package: autogram labels Sep 18, 2025
@codecov
Copy link
Copy Markdown

codecov bot commented Sep 18, 2025

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines Coverage Δ
src/torchjd/autogram/_engine.py 100.00% <100.00%> (ø)
src/torchjd/autogram/_module_hook_manager.py 100.00% <100.00%> (ø)
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Copy link
Copy Markdown
Contributor

@PierreQuinton PierreQuinton left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM if typing is prohibitively annoying.

@PierreQuinton PierreQuinton merged commit f5a8996 into main Sep 19, 2025
17 checks passed
@PierreQuinton PierreQuinton deleted the fix-autogram-ref-cycles branch September 19, 2025 14:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cc: fix Conventional commit type for bug fixes of the actual library (changes to src). package: autogram

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Reference cycles in autogram

2 participants