Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adoption cache should use WeakValueDictionary. #2495

Merged
merged 1 commit into from
Oct 3, 2022

Conversation

levskaya
Copy link
Collaborator

@levskaya levskaya commented Oct 2, 2022

We use a cache to preserve sharing-by-reference relationships when "adopting" outside modules into a new top-level functional module context (adoption occurs frequently when using advanced config systems like gin or fiddle, or during interactive development).

Previously we were careful not to hold references to Scope objects by having the cache itself be a WeakKeyDictionary, however the scope-keyed entries were just classic dicts, and this was leading to leaks of tracers and arrays ultimately due to long-lived references to submodule objects from the cache. By using a WeakValueDictionary to refer to these we avoid these lifetime issues.

We use a cache to preserve sharing-by-reference relationships when
"adopting" outside modules into a new top-level functional module
context (adoption occurs frequently when using advanced config systems
like gin or fiddle, or during interactive development).

Previously we were careful not to hold references to Scope objects by
having the cache itself be a WeakKeyDictionary, however the scope-keyed
entries were just classic dicts, and this was leading to leaks of
tracers and arrays ultimately due to long-lived references to submodule
objects.  By using a WeakValueDictionary to refer to these we avoid
these lifetime issues.
@levskaya
Copy link
Collaborator Author

levskaya commented Oct 2, 2022

Fixes #2493

@levskaya levskaya requested a review from jheek October 2, 2022 18:54
@codecov-commenter
Copy link

codecov-commenter commented Oct 2, 2022

Codecov Report

Merging #2495 (53a4862) into main (ef67f1c) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main    #2495   +/-   ##
=======================================
  Coverage   78.87%   78.88%           
=======================================
  Files          48       48           
  Lines        5118     5120    +2     
=======================================
+ Hits         4037     4039    +2     
  Misses       1081     1081           
Impacted Files Coverage Δ
flax/linen/module.py 92.80% <100.00%> (+0.02%) ⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

Copy link
Member

@jheek jheek left a comment

Choose a reason for hiding this comment

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

Thanks! Makes sense

@copybara-service copybara-service bot merged commit 2650849 into google:main Oct 3, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants