Adding random number generation API

[jakirkham]

In conjunction with the Array API, often need some form of random number generation as well. Since this would generate Arrays and be used with Arrays the user has, there is some benefit from downstream users perspective if there is a standard API for handling random number generation.

For example in scikit-learn, the following discussion ( scikit-learn/scikit-learn#22352 (comment) ) may shed some light. A couple points that came up were whether to have a state or stateless API and NumPy's old and new APIs.

[leofang]

A couple points that came up were whether to have a state or stateless API and NumPy's old and new APIs.

I'd like to quote @rgommers from there (scikit-learn/scikit-learn#22352 (comment)) because I feel it clarified a lot:

It's not a stateless PRNG (there's no such thing)

[jakirkham]

Yeah was just summarizing the discussion up to that point.

Stateless isn't really the right term for this I think. A more correct description might be functional as many functional languages use a similar strategy. IOW something like this

state0 = seed(42)
state1, n = randint(state0)
state2, m = randint(state1)

IOW the user is responsible for tracking the state in this model and threading it through subsequent random number generation function calls.

The other model is object oriented like NumPy's where state lives in some object that gets mutated each time a random number is generated.

[rgommers]

We had a previous discussion on PRNG APIs, and IIRC no one was attached to the legacy numpy.random API (which most libraries currently have as their default - although TensorFlow is quite different: tf.random.Generator, also class-based).

The new NumPy API seems a little more user-friendly than the JAX APIs, and offers more functionality too. And it's no less safe - I think it's just as well possible to shoot yourself in the foot with JAX: just forget to manually create a new sub-key once. This is more of a philosophical difference: there's only one way of doing things in JAX, which is more verbose but the same for serial and parallel execution - while NumPy has a more concise way for serial (its default mode), but the user must remember to use a second method (equivalent to jax.random.split) to get the right parallel behavior.

One advantage of the JAX API is that there's only one way to do things, while NumPy has multiple ways to deal with parallelism - so more functionality, but also harder to understand or standardize. That may be partly due to providing multiple PRNG algorithms though, while JAX provides only the Threefry algorithm (not entirely true, there's a second experimental one - XLA Random Bit Generator). PyTorch for example provides Philox (which JAX also could have used) and MT19937. Now a standard doesn't have to deal with exact reproducibility across libraries, but it should allow a design that allows libraries to choose their own algorithms. I think this works with either JAX's key (which contains an algorithm selector), or NumPy's class-based API. NumPy's SeedSequence is like the seeding part of JAX's key, and should work with any algorithm too. NumPy's .jumped cannot be supported by all algorithms AFAIK.

The main issues for standardization I see are:

• (a) JAX cannot adopt any class-based API,
• (b) NumPy probably isn't going to be very receptive to yet another API, and
• (c) it'd be a lot of work for other libraries like PyTorch to switch from their old API to any other API.

tl;dr this will not be easy to standardize

[jakirkham]

Thanks for the detailed write up Ralf! 🙏

FWIW Dask is already working on adopting the new NumPy API ( dask/dask#9038 )

It looks like CuPy already did this in 9.0.0 with PR ( cupy/cupy#4177 ) (though Leo should feel free to correct me)

That all being said, maybe it is worth asking if a subset of NumPy's API might be easier to adopt and if so what that looks like.

Also the other important question here is what API is going to be most useful for downstream users. Started with scikit-learn as they make a lot of use of random number generation and creating some usable API with scikit-learn would be a win. Though maybe there are other downstream libraries that would make sense as well (perhaps statsmodels? others?).

[oleksandr-pavlyk]

@leofang While there is no stateless PRNG, HW-based true random number generators have empty state (MKL's NONDETERM basic random number generator) and these should be supported by the spec as well.

[rgommers]

To get a better feel for the tradeoffs, here is code for a couple of things:

1. Serial and parallel (multiprocessing) usage for numpy.random and jax.random
2. An implementation of the JAX API (PRNGKey, split, uniform) on top of the numpy.random infrastructure. Vice-versa is not possible.

For all the below code in a single gist, see here.

import secrets
import multiprocessing

import numpy as np
import jax


USE_FIXED_SEED = False

if USE_FIXED_SEED:
    seed = 38968222334307
else:
    # Generate a random high-entropy seed for use in the below examples
    # jax.random.PRNGKey doesn't accept None to do this automatically
    seed = secrets.randbits(32)  # JAX can't deal with >32-bits


# NumPy serial
rng = np.random.default_rng(seed=seed)
vals = rng.uniform(size=3)
val = rng.uniform(size=1)

# NumPy parallel
sseq = np.random.SeedSequence(entropy=seed)
child_seeds = sseq.spawn(4)
rngs = [np.random.default_rng(seed=s) for s in child_seeds]

def use_rngs_numpy(rng):
    vals = rng.uniform(size=3)
    val = rng.uniform(size=1)
    print(vals, val)

def main_numpy():
    with multiprocessing.Pool(processes=4) as pool:
        pool.map(use_rngs_numpy, rngs)


# JAX serial (also auto-parallelizes fine by design)
key = jax.random.PRNGKey(seed)
key, subkey = jax.random.split(key)  # this one could be left out, but best practice is probably to always use `split` first
vals = jax.random.uniform(subkey, shape=(3,))
key, subkey = jax.random.split(key)  # don't forget this!
val = jax.random.uniform(subkey, shape=(1,))


# JAX parallel with multiprocessing
def use_rngs_jax(key):
    key, subkey = jax.random.split(key)
    vals = jax.random.uniform(subkey, shape=(3,))
    key, subkey = jax.random.split(key)
    val = jax.random.uniform(subkey, shape=(1,))
    print(vals, val)


def main_jax():
    key = jax.random.PRNGKey(seed)
    key, *subkeys = jax.random.split(key, 5)  # gotcha: "5" gives us 4 subkeys
    with multiprocessing.Pool(processes=4) as pool:
        pool.map(use_rngs_jax, subkeys)


if __name__ == '__main__':
    # JAX does not work with the default `fork` (due to internal threading)
    multiprocessing.set_start_method('forkserver')

    print('\nNumPy with multiprocessing:\n')
    main_numpy()
    print('\n\nJAX with multiprocessing:\n')
    main_jax()

JAX does seem to have a few gotchas with seed creation, it can't deal with high-entropy seeds apparently (at least in 0.2.27, released 18 Jan 2022):

In [24]: seed = secrets.randbits(64)

In [25]: jax.random.PRNGKey(seed)
---------------------------------------------------------------------------
OverflowError                             Traceback (most recent call last)
<ipython-input-25-7a8d328c270c> in <module>
----> 1 jax.random.PRNGKey(seed)

~/anaconda3/envs/many-libs/lib/python3.7/site-packages/jax/_src/random.py in PRNGKey(seed)
     57   # Explicitly cast to int64 for JIT invariance of behavior on large ints.
     58   if isinstance(seed, int):
---> 59     seed = np.int64(seed)
     60   # Converting to jnp.array may truncate bits when jax_enable_x64=False, but this
     61   # is necessary for the sake of JIT invariance of the result for such values.

OverflowError: Python int too large to convert to C long

In [26]: seed = secrets.randbits(128)

In [27]: jax.random.PRNGKey(seed)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-27-7a8d328c270c> in <module>
----> 1 jax.random.PRNGKey(seed)

~/anaconda3/envs/many-libs/lib/python3.7/site-packages/jax/_src/random.py in PRNGKey(seed)
     53     raise TypeError(f"PRNGKey seed must be a scalar; got {seed!r}.")
     54   if not np.issubdtype(np.result_type(seed), np.integer):
---> 55     raise TypeError(f"PRNGKey seed must be an integer; got {seed!r}")
     56
     57   # Explicitly cast to int64 for JIT invariance of behavior on large ints.

TypeError: PRNGKey seed must be an integer; got 67681183633192462759155065893448052088

In [28]: seed = secrets.randbits(64)

In [29]: jax.random.PRNGKey(seed)
---------------------------------------------------------------------------
OverflowError                             Traceback (most recent call last)
<ipython-input-29-7a8d328c270c> in <module>
----> 1 jax.random.PRNGKey(seed)

~/anaconda3/envs/many-libs/lib/python3.7/site-packages/jax/_src/random.py in PRNGKey(seed)
     57   # Explicitly cast to int64 for JIT invariance of behavior on large ints.
     58   if isinstance(seed, int):
---> 59     seed = np.int64(seed)
     60   # Converting to jnp.array may truncate bits when jax_enable_x64=False, but this
     61   # is necessary for the sake of JIT invariance of the result for such values.

OverflowError: Python int too large to convert to C long

In [30]: seed = secrets.randbits(32)

In [31]: jax.random.PRNGKey(seed)
Out[31]: DeviceArray([         0, 3279739543], dtype=uint32)

A JAX-style API with numpy.random under the hood:

"""
Implement `jax.random` APIs with NumPy, and `numpy.random` APIs with JAX.
The purpose of this is to be able to compare APIs more easily, and clarify
where they are and aren't similar.
"""

import secrets
import multiprocessing

import numpy as np
import jax


USE_FIXED_SEED = False

if USE_FIXED_SEED:
    seed = 38968222334307
else:
    # Generate a random high-entropy seed for use in the below examples
    # jax.random.PRNGKey doesn't accept None to do this automatically
    seed = secrets.randbits(32)  # JAX can't deal with >32-bits



def PRNGKey(seed):
    """
    Create a key from a seed. `seed` must be a 32-bit (or 64-bit?) integer.
    """
    # Note: selecting a non-default PRNG algorithm is done via a global config
    #       flag (not good, should be a keyword or similar ...)
    seed = np.random.SeedSequence(seed)
    rng = np.random.default_rng(seed)
    key = (seed, rng)
    return key


def split(key, num=2):
    """
    Parameters
    ----------
    key : tuple
        Size-2 tuple, the first element a `SeedSequence` instance, the second
        containing the algorithm selector.
    num : int, optional
        The number of keys to produce (default: 2).

    Returns
    -------
    keys : tuple of 2-tuples
        `num` number of keys (each key being a 2-tuple)
    """
    seed, rng = key
    child_seeds = seed.spawn(num)
    keys = ((s, rng) for s in child_seeds)
    return keys


def uniform(key, shape=(), dtype=np.float64, minval=0.0, maxval=1.0):
    seed, rng = key
    # Creating a new Generator instance from an old one with the same
    # underlying BitGenerator type requires using non-public API:
    rng = np.random.Generator(rng._bit_generator.__class__(seed))
    return rng.uniform(low=minval, high=maxval, size=shape).astype(dtype)


def use_jaxlike_api(key=None):
    if key is None:
        key = PRNGKey(seed)

    key, subkey = split(key)
    vals = uniform(subkey, shape=(3,))
    key, subkey = split(key)  # don't forget this!
    val = uniform(subkey, shape=(1,))
    print(vals, val)


def use_jaxlike_api_mp():
    key = PRNGKey(seed)
    key, *subkeys = split(key, 5)
    with multiprocessing.Pool(processes=4) as pool:
        pool.map(use_jaxlike_api, subkeys)


if __name__ == '__main__':
    # JAX does not work with the default `fork` (due to internal threading)
    multiprocessing.set_start_method('forkserver')

    print('\n\nUse JAX-like API (serial):\n')
    use_jaxlike_api()
    print('\n\nUse JAX-like API (multiprocessing):\n')
    use_jaxlike_api_mp()

A couple of thoughts:

• JAX has several rough edges, like not being able to deal with all Python integers, not working well with multiprocessing, not allowing seed=None, and the num keyword to split being odd (to use 4 subkeys in your parallel generation, you need to say num=5).
  • The first one and last two are fairly minor.
  • The multiprocessing one is quite annoying, but an inherent issue due to JAX's and multiprocessing's internals not getting along, and in principle unrelated to API design.
• Tradeoff: NumPy's API is more concise for both serial usage and parallel usage with multiprocessing. However, JAX's API works naturally with auto-parallelization/threading, while with NumPy the user has to switch to a completely different method using SeedSequence (note: .jumped is similar, doesn't change this tradeoff).
  • How do you shoot yourself in the foot with JAX? Forget to use split() before drawing random numbers.
  • How do you shoot yourself in the foot with NumPy? For serial use (NumPy's default): safe, no footguns. For parallel use: quite easy to shoot yourself in the foot, users have to realize that they need to switch paradigm and then fiddle with SeedSequence or .jumped.
• Creating new BitGenerator and Generator instances each time one wants to generate more random numbers (see the uniform implementation) will have significant overhead. I haven't measured, but I'm fairly sure it matters. This can probably not be avoided with the current design, would need a larger change in numpy.random to make it efficient I think (disclaimer: I haven't thought about this super hard).

One other thing to point out: the JAX docs comparing to NumPy are wildly outdated/unfair, they use the non-recommended (global state) way of using the legacy API. In general, numpy.random provides a superset of capabilities of those of jax.random, as evidenced by how easy it is to implement a JAX-like API on top of the current numpy.random. Related: see this informative comment by @rkern for more thoughts on jax.random vs. numpy.random. In particular this comment:

Just to provide some mathematical background, Jax's PRNG is in the same weak-crypto family as our Philox BitGenerator. The method by which it splits only works well for that family of weak-crypto PRNGs because that family keeps its initial seed around as the key value and only evolves a counter as one draws numbers from it. The other PRNGs iterate the state.

My tentative conclusions based on the above:

• If consumer libraries want to write portable code across NumPy, JAX, PyTorch and other libraries with NumPy-matching APIs (Dask, CuPy), then there's two options:
  • (1) A shim like the above JAX-like API.
  • (2) Use the numpy.random API as the common case, and simply special-case JAX.
  • It will be quite an effort to teach everyone (1), however I'm not sure that's a show-stopping concern since the new-style numpy.random.Generator API is still not used much (even experienced scikit-learn devs weren't aware of it in a recent thread). I've given two talks that covered the new numpy.random.Generator API, and each time many people were surprised (and enthusiastic) - so there's low uptake. Which would be nice to improve on either way.
  • That said, (2) is going to be more performant, given that split only works well for counter-based algorithms, and special-casing JAX is not that much work.
• The JAX API needs some improvements for the rough edges, but is in principle easier to teach.
• Regarding @jakirkham's question "That all being said, maybe it is worth asking if a subset of NumPy's API might be easier to adopt and if so what that looks like.": I think SeedSequence is probably the thing to include for parallelism; having multiple ways of doing parallelism is nice for power users, but probably confusing to many others.

This is all a little nontrivial, so let me ping a few folks for input: @rkern for design/implementation thoughts and whether I missed anything important related to the NumPy implementation. @shoyer, @apaszke for thoughts from the JAX side.

[rkern]

• For parallel use: quite easy to shoot yourself in the foot, users have to realize that they need to switch paradigm and then fiddle with SeedSequence or .jumped.

I'll need to spend more time to read the whole thread, but I will add here that we have always had a plan to lift .spawn() up to Generator. We left it out in the initial release to build some more comfort with SeedSequence spawning before we committed to it. But I think it's time. It would make it a lot easier to avoid the foot-guns without having to go back and change a lot of code to pass around more state.

[rkern]

Specifically, this code:

# NumPy parallel
sseq = np.random.SeedSequence(entropy=seed)
child_seeds = sseq.spawn(4)
rngs = [np.random.default_rng(seed=s) for s in child_seeds]

would become:

rngs = rng.spawn(4)

[rgommers]

I'll need to spend more time to read the whole thread, but I will add here that we have always had a plan to lift .spawn() up to Generator.

Thanks - I think that would indeed be quite helpful!

[jakirkham]

Sorry for a bit of a tangent here, but is calling .spawn() incrementally valid? Like

rng1 = rng.spawn(1)
rng2 = rng.spawn(1)
...
rngn = rng.spawn(1)

This can come up when new children processes are created/destroyed on-demand (IOW autoscaling).

[seberg]

Yes it is. More importantly:

rng3 = rng2.spawn()

is also valid. You can spawn any children to get more independent streams. There will never be a collision (at least not within reasonable probabilities).

[rgommers]

I went searching through some older notes and found this from @alextp: "I think the direction in TensorFlow is to follow JAX. Have functions for stateless random generation. They compose well. But they are not ergonomic. Layer on top of this will be something stateful. Once you do this, you introduce checkpoints etc, for determinism. So I would OK to add stateless API, but wouldn't be okay to add stateful one." That's from more than a year ago, so I don't know if it has been implemented like that, or anything changed in the meantime in TensorFlow. Maybe @edloper you can tell us?

[rkern]

I'm perfectly content with having 0 PRNG APIs in the standard (and far prefer 0 in the standard to having 2 in the standard). It seems like there is significant variety in what different communities need and want.

[shoyer]

I think it's equally straightforward to write an explicitly stateful RNG system like numpy.random.Generator using JAX. Here's a prototype:

import jax


class JaxGenerator:
  def __init__(self, state):
    self.state = state

  def uniform(self):
    self.state, key = jax.random.split(self.state)
    return jax.random.uniform(key)
  
  def spawn(self, count):
    self.state, *keys = jax.random.split(self.state, count + 1)
    return [JaxGenerator(key) for key in keys]

  def __repr__(self):
    return f'{type(self).__name__}(state={self.state})'


def jax_default_rng(seed):
  return JaxGenerator(jax.random.PRNGKey(seed))


rng = jax_default_rng(0)
print(rng.uniform())  # 0.10536897
print(rng.uniform())  # 0.2787192
print(rng)  # JaxGenerator(state=[2384771982 3928867769])

rng2 = JaxGenerator(rng.state)
print(rng2)  # JaxGenerator(state=[2384771982 3928867769])

rngs = rng.spawn(2)
print(rngs)  # [JaxGenerator(state=[1777981902 3244208681]), JaxGenerator(state=[ 669635267 2816531647])]

From a JAX design perspective, stateful RNGs like this are not encouraged, because JAX's function transforms will break if you pass stateful objects into them. But you can still safely use this sort of API with JAX, as long as you're careful to create Generator objects inside pure functions, e.g.,

@jax.jit
@jax.vmap
def batched_random_uniform(seed):
  return jax_default_rng(seed).uniform()

print(batched_random_uniform(jnp.arange(5)))
# DeviceArray([0.10536897, 0.12568676, 0.4336251 , 0.47652578, 0.7844808 ],            dtype=float32)

Or explicitly keeping track of updated RNG state:

@jax.jit
def explicit_state_random_uniform(state):
  rng = JaxGenerator(state)
  sample = rng.uniform()  # must happen *before* calculating new_state 
  new_state = rng.state
  return new_state, sample

state = jax.random.PRNGKey(0)
state2, sample = explicit_state_random_uniform(state)
print(state2, sample)

In fact, Haiku, which is one of the most popular neural net libraries in JAX, does something very similar with haiku.next_rng_key.

These stateful interfaces are perhaps easier to misuse with JAX's functional transforms than pure functions, but are still a major improvement over global state.