Fix PRNG handling in `nn.jit` under `nn.scan`.

[copybara-service[bot]]

Fix PRNG handling in nn.jit under nn.scan.

• 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.

[google-cla[bot]]

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.