Experiment: mutable RNG object

[jakevdp]

An experiment with implicitly-updated PRNG keys in JAX, based on the MutableArray experiment. The API is inspired by numpy.random.default_rng.

The mutable array code is still in development, with a few blockers to being able to rely on it for non-experimental APIs, but I thought this would be a useful exploration in the meantime!

Example:

In [1]: from jax.experimental.mutable_rng import default_rng

In [2]: rng = default_rng(123456)

# Calling rng.key() repeatedly will generate new keys each time, implicitly updating
# the internal state:
In [3]: rng.key()
Out[3]: 
Array((), dtype=key<fry>) overlaying:
[2376306259 4203193672]

In [4]: rng.key()
Out[4]: 
Array((), dtype=key<fry>) overlaying:
[ 989706967 1551494763]

# The rng object provides methods to generate values based on these implicitly-updated
# keys, inspired by the api of np.random.default_rng()
In [5]: rng.uniform()
Out[5]: Array(0.6851474, dtype=float32)

In [6]: rng.uniform()
Out[6]: Array(0.70181704, dtype=float32)

In [7]: rng.integers(0, 10, size=5)
Out[7]: Array([8, 0, 5, 7, 2], dtype=int32)

In [8]: rng.integers(0, 10, size=5)
Out[8]: Array([9, 0, 6, 7, 1], dtype=int32)

# Even when used under JIT, the implicit update propagates through due to being
# backed by the experimental MutableArray object.
In [9]: import jax

In [10]: @jax.jit
    ...: def f(rng):
    ...:   return rng.integers(0, 10, size=5)
    ...: 

In [11]: f(rng)
Out[11]: Array([7, 1, 0, 8, 1], dtype=int32)

In [12]: f(rng)
Out[12]: Array([3, 9, 2, 5, 6], dtype=int32)