rnglib
------------------------

.. automodule:: flax.nnx
.. currentmodule:: flax.nnx

.. autoclass:: Rngs
   :members: __init__
.. autoclass:: RngStream
   :members:
.. autofunction:: split_rngs
.. autofunction:: fork_rngs
.. autofunction:: reseed
.. autofunction:: with_rngs