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

Fix setup trigger issue under sharing and transforms. #3153

Merged
merged 1 commit into from
Jun 20, 2023

Conversation

levskaya
Copy link
Collaborator

There was a subtle bug in the case of using a shared module that defined a setup() method that was called first at one location under a transform (like remat or jit) and at another location without the transform.

In the latter case the "outside" instance of the submodule had had a "shallow" _try_setup call pre-transform and was then marked as having had setup called, but the setup method had not in fact been called. This requires a one-line logic fix in _try_setup().

There was a subtle bug in the case of using a shared module that defined
a setup() method that was called first at one location under a transform
(like remat or jit) and at another location without the transform.

In the latter case the "outside" instance of the submodule had had a
"shallow" _try_setup call pre-transform and was then marked as having
had setup called, but the setup method had not in fact been called.
This requires a one-line logic fix in _try_setup().
@codecov-commenter
Copy link

codecov-commenter commented Jun 20, 2023

Codecov Report

Merging #3153 (5a82a05) into main (fcd914d) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main    #3153   +/-   ##
=======================================
  Coverage   82.12%   82.12%           
=======================================
  Files          55       55           
  Lines        6063     6064    +1     
=======================================
+ Hits         4979     4980    +1     
  Misses       1084     1084           
Impacted Files Coverage Δ
flax/linen/module.py 92.82% <100.00%> (+<0.01%) ⬆️

Copy link
Collaborator

@cgarciae cgarciae left a comment

Choose a reason for hiding this comment

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

Looks good. Sharing + transforms is hard to reason about.

@copybara-service copybara-service bot merged commit 875b40f into google:main Jun 20, 2023
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