flax
139d02fd - Fix initialization RNGs to work with omnistaging for jitted inits.

Commit

5 years ago

Fix initialization RNGs to work with omnistaging for jitted inits. PiperOrigin-RevId: 332326041

References

Author

levskaya

Committer

a-googler

Parents