Give RngStream the same interface as Jax's stateful RNGs
Open
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
Jax will likely support a stateful PRNG object similar to Flax's RngStreams in the future: https://github.com/jax-ml/jax/pull/28845.
Eventually, we'll want to have flax use these native flax objects. To ease that eventual transition, we should give RngStreams the same interface (specifically, the `key` and `split` methods) and express the new explicit-rng docs in terms of this interface.
Contributor guide
Assessment
This issue has not been assessed yet.