Fix PRNG handling in nn.jit under nn.scan.
#4359
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
![copybara-service[bot]](https://avatars.githubusercontent.com/in/44061?s=60&v=4){kind=small}
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
copybara-service
bot
force-pushed
the
test_692231015
branch
from
8c9e450
to
b004ce7
Compare
copybara-service
bot
changed the title
nn.jit under nn.scan.nn.jit under nn.scan.
copybara-service
bot
force-pushed
the
test_692231015
branch
from
b004ce7
to
97e19f3
Compare
* `nn.scan` does an abstract eval before compilation to check for constants that are then traced out. Before this change, the abstract eval increments static RNG counters, which creates a side-effect where RNG counters are not properly updated once inner functions are jitted (i.e., under `nn.jit`). * In this fix, we cache the impact a first pass through `nn.scan` and `nn.jit` would have on rng counters and "replay" that impact on subsequent passes so that rng state remains unaffected. * This solution doesn't affect PRNG derivations outside this isolated case and is a placeholder while a more permanent solution, which would affect PRNG derivations, is worked out. PiperOrigin-RevId: 694463650
copybara-service
bot
force-pushed
the
test_692231015
branch
from
97e19f3
to
0f631a2
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Fix PRNG handling in
nn.jitundernn.scan.nn.scandoes an abstract eval before compilation to check for constants that are then traced out. Before this change, the abstract eval increments static RNG counters, which creates a side-effect where RNG counters are not properly updated once inner functions are jitted (i.e., undernn.jit).nn.scanandnn.jitwould have on rng counters and "replay" that impact on subsequent passes so that rng state remains unaffected.