Adding random number generation API

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.

@rgommers

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)

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.

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

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?).

@leofang

@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.

@jakirkham

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

Serial and parallel (multiprocessing) usage for numpy.random and jax.random
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):

Details

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.

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.

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)

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!

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).

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).

@alextp

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?

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.

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.

To bridge the gap between NumPy's and JAX's random number APIs, I would suggest slightly extending NumPy's API so it's easier to be explicit about state. Namely, we should add the spawn() method, and also support directly accessing/setting .state on Generator objects, rather than only on BitGenerator.

Here's my minimal wrapper of NumPy's RNG API with these extensions, to match the JAX API above:

import numpy as np


class NumpyGenerator:

  def __init__(self, state):
    # TODO: avoid initializing this dummy Generator/BitGenerator?
    self._rng = np.random.default_rng()
    self.state = state

  def uniform(self):
    return self._rng.uniform()
  
  def spawn(self, count):
    entropy = self._rng.integers(2**63)  # TODO: better entropy?
    sseq = np.random.SeedSequence(entropy)
    child_seeds = sseq.spawn(count)
    return [numpy_default_rng(seed) for seed in child_seeds]

  @property
  def state(self):
    return self._rng.bit_generator.state

  @state.setter
  def state(self, value):
    self._rng.bit_generator.state = value

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


def numpy_default_rng(seed):
  return NumpyGenerator(np.random.default_rng(seed).bit_generator.state)


rng = numpy_default_rng(0)
print(rng.uniform())  # 0.6369616873214543
print(rng.uniform())  # 0.2697867137638703
print(rng)  # NumpyGenerator(state={'bit_generator': 'PCG64', 'state': {'state': 143609658456486183636066271097634410721, 'inc': 87136372517582989555478159403783844777}, 'has_uint32': 0, 'uinteger': 0})

rng2 = NumpyGenerator(rng.state)
print(rng2)  # NumpyGenerator(state={'bit_generator': 'PCG64', 'state': {'state': 143609658456486183636066271097634410721, 'inc': 87136372517582989555478159403783844777}, 'has_uint32': 0, 'uinteger': 0})

rngs = rng.spawn(2)
print(rngs)  # [NumpyGenerator(state={'bit_generator': 'PCG64', 'state': {'state': 164868448360684748498847325894109072011, 'inc': 241524822143570234404080558697197945801}, 'has_uint32': 0, 'uinteger': 0}), NumpyGenerator(state={'bit_generator': 'PCG64', 'state': {'state': 219573644955839246335449654370252341036, 'inc': 175512849395095609630857841553467033115}, 'has_uint32': 0, 'uinteger': 0})]

Thanks for the pointers - in that case I think I will link to this thread, because it seems more direct than linking to a long doc and pointing out the omission of the topic in question.

Are there actually large reasons why the NumPy API (with .spawn() also living on the rng object itself and a default_rng() function), is not a potential way forward?

It seems that there are the 3 "ways" of using a PRNG that Robert listed also above?:

Sequential drawing.
"Lazy" sequential drawing (rng, data = uniform(rng, size=1000) would be a "stateless" spelling of that. I.e. a way that allows data to be evaluated at a later time, but avoids splitting.)
- For an end-user/library API (i.e. what we are discussing) it may not matter if the provider (JAX) actually has to do splitting to achieve this?
Flexible splitting/spawning for parallelism, that does not come with the limits of 2.
- A constraint here is maybe that NumPy guarantees that splitting does not depend on the RNG state as modified by sequential draws. But I expect that this would be easy to guarantee in all cases (if desired).

Everyone provides 3, some PRNGs (as NumPy) can probably not provide 2 but don't need it (NumPy never evaluates lazily – it might only be a curious addition for parallelization). JAX need 2 (but is currently "missing" it – i.e. is using the less optimal "always spawn" scheme).

But I don't think it has to be user-facing API? If you write:

data1 = rng.uniform()
data2 = rng.uniform()

in JAX or NumPy, whether that draws the numbers directly sequentially, or does so lazily (with some advancing or splitting scheme) hardly matters? What matters is that the user facing API is "sequential", because writing it differently would be bad for many PRNG schemes.

What I don't quite understand yet, is why a stateful API that provides:

Sequential draws as the normal way to get random numbers
A way to spawn off new streams for parallelism

is not a reasonable start. It would be bad for NumPy to provide a JAX-style stateless API, but why can't JAX provide a sequential stateful API (even if it is not necessarily a sequential RNG generation internally).
After all, just because the API is "sequential", does not mean it has to guarantee that:

concat([rng.uniform(size=10), rng.uniform(size= 20)])

gives the same as:

concate([rng.uniform(size= 20), rng.uniform(size= 10)])

NumPy will do this (often or always?) but it doesn't seem necessary for the end-user API?

The difficulty I see a bit is that some implementations may look like they provide guarantees like the above concatenation and others will not provide those same guarantees.

@rkern

It's a welcome surprise to have this thorough review of the JAX PRNG! Thanks @rkern for the really nice detailed look in particular.

I'm only arriving to this thread now and still catching up, but as a quick note for now: we happen to be actively working on some changes to the threefry_random_bits implementation, most immediately in how it lays out the counter values from iota, with the aim of making it more friendly to partitioning and to JAX's experimental support of dynamic array shapes. I don't think the immediate-term changes will affect many of @rkern's overall takeaways, but they may address this remark in part:

In particular, instead of incrementing the counter in 64-bit blocks, it creates a 32-bit counter array, splits it in half to use the first half as the upper 32-bit word and the second half as the lower 32-bit word. That's a bit wasteful (you could instead just use zeros(rem//2) and lax.iota(rem//2) for rem requested words instead), but they are valid inputs. The two arrays which are output (being separately the upper and lower words of the cipher_text blocks) are then concatenated together end-to-end instead of reassembling those words. I assume this is done because it's cheaper to concatenate than to interleave again. This is fine as far as it goes; it doesn't hugely affect the quality of numbers coming out of one jax.random.uniform() call, for example. But it does wreck the bijection properties, which will become important when we talk about implementation choices for split().

We indeed do not need to lay out the counter values exactly in the way that we do, so you'll see changes to that end. (I've observed in other contexts over time that the particular split/concat scheme we use isn't ideal, but I only picked up work on this again last week.)

So far, so good. Despite that quirk, I have only one qualm with the core PRNG scheme for drawing bits in the jax.random functions: 64 bits is quite small. I consider 128-bit PRNGs to be de rigeur today, particularly in contexts where we are talking about large amounts of parallelism (and the JAX API enforces a large amount of parallelism).

As of recently, I'd say that the key size—and in fact the entire hash function or base generator—is not an essential or timeless choice for us. It is what we implemented initially, and it remains the default for now. But as of jax-ml/jax#6899 we have a means of replacing the underlying bit generator. We use this internally to experiment with and offer other bit generators. I opened jax-ml/jax#7676 a while back precisely to track the introduction of a 128-bit generator. We also have plans to make it possible for users to plug in an arbitrary generator of their own, and to involve differently-backed PRNGs in the same process (all tracked at jax-ml/jax#9263). We might always choose to change our default away from the current threefry2x32 hash at some point.

So, where possible in the current discussion, it may be useful to assume that the base generator can change to meet the needs of the machine, application, etc.

@froystig

@froystig That's all good to hear! Like I said in my introduction, I didn't consider any of those details as immutable black marks against JAX or the overall design, but they did serve as a way to talk about the similarities with SeedSequence, the kinds of improvements that are left in the design space (bijections as far as the eye can see!), and what limits are still fundamental to a splitting-centric design (the Birthday Bound). And I'm a long-winded jerk.

I don't know all of JAX's rewriting capabilities, but I don't think it can handle whatever it needs to do to make the Generator-like API work with the same reordering benefits.

data1 = rng.uniform()
data2 = rng.uniform()

What it would have to do here is recognize that rng is a Generator, so the special case applies, and magically under the covers transform the code to do a couple of splits up at the top before any of the uniform() calls.

And even if JAX did have technical capability of doing so, I'm not sure that they want their users to mix mental models like that. It's hard enough to teach people to use one consistently.

@rkern

they did serve as a way to talk about the similarities with SeedSequence, the kinds of improvements that are left in the design space (bijections as far as the eye can see!), and what limits are still fundamental to a splitting-centric design (the Birthday Bound).

@rkern Absolutely – I understood that and it makes sense. I only meant to share some related thoughts and work that's in progress, mostly as an aside. Your comments are super valuable, and this thread is a rare opportunity for us on JAX. Thank you!

I'm the author of tf.random.Generator (link) and several docs (e.g. this and this). TF RNG's current status is using stateless RNGs (similar to JAX's) as the underlying engine (code), and a mutable counter (implemented as a tf.Variable) to advance the RNG state (code). tf.random.Generator is basically stateless RNGs + a mutable counter. It works well in TF because TF is more friendly to mutability than JAX.

Unlike JAX, tf.random.Generator supports but doesn't do splitting unless asked to (code). It just does counter increment in common (sequential) cases (code). I feel splitting is more wasteful of the (key, counter) space than counter increment, but I'm not sure. Counter increment is of course computationally cheaper than splitting, the latter of which caused some speed problems in some TF models.

TF only supports counter-based RNG algorithms and has no plan to support non-counter-based ones.

Yes, I can confirm that splitting is more profligate of the state space than the counter increment, especially when you constrain yourself to 64-bit keys, if I'm reading the code correctly. I think the same math from my analysis of JAX's PRNG splitting applies (the only real difference is that tf.random.Generator will generate the new keys at whatever counter it is currently on rather than JAX always starting at 0, but that doesn't affect the non-invertible mapping argument).

From a user perspective, specifically someone writing libraries which want to support arbitrary array types, I don't see the problem with carrying the random state separately from the random namespace. It is not the slickest interface, but it can be implemented today:

# file: random_numpy.np

def standard_normal(rng, shape=()):
    return rng, rng.standard_normal(size=shape)

def poisson(rng, lam, shape=()):
    return rng, rng.poisson(lam, size=shape)


# file: random_jax.py

from jax import random

def standard_normal(key, shape=()):
    key, subkey = random.split(key)
    return key, random.normal(subkey, shape)

def poisson(key, lam, shape=()):
    key, subkey = random.split(key)
    return key, random.poisson(subkey, lam, shape)


# file: test.py

def random_namespace(rs):
    import sys
    if 'numpy' in sys.modules:
        import numpy as np
        if isinstance(rs, np.random.Generator):
            return __import__('random_numpy')
    if 'jax' in sys.modules:
        import jax
        if isinstance(rs, jax.Array):
            return __import__('random_jax')

def my_function(lam, rs):
    random = random_namespace(rs)
    rs, rv1 = random.standard_normal(rs, shape=(4,))
    rs, rv2 = random.poisson(rs, lam, shape=(4,))
    return rv1 + rv2

import numpy as np
rng = np.random.default_rng()
print(my_function(1.0, rng))
# [ 1.25035394 -0.11511349  1.87203598  2.55088409]

import jax
key = jax.random.PRNGKey(42)
print(my_function(1.0, key))
# [0.43244982 0.8856307  0.06793922 1.4641162 ]

For anything more complicated than sampling random variates, particularly anything that requires knowledge of the implementation details of the random state, we accept that we will always have to handle each framework separately.

The main difference is that rng in numpy and key in JAX are fundamentally different kinds of objects. In the numpy example, the call to my_function has the side-effect of mutating rng in-place, so that subsequent calls will return different random values. In JAX, functions are pure, and so the call to my_function cannot mutate key, and subsequent calls with the same key will produce identical values to the first call:

import numpy as np
rng = np.random.default_rng()
print(my_function(1.0, rng))
# [ 1.25035394 -0.11511349  1.87203598  2.55088409]
print(my_function(1.0, rng))
# [3.15990919 1.07582056 1.09202392 0.33987543]

import jax
key = jax.random.PRNGKey(42)
print(my_function(1.0, key))
# [-0.5675502   1.8856307   0.06793922  3.464116  ]
print(my_function(1.0, key))
# [-0.5675502   1.8856307   0.06793922  3.464116  ]

Yes, there is probably no universal API for situations where it matters whether or not we are dealing with a stateful RNG or a stateless key.

My point is that it does not preclude us from having a universal (functional) API that can generate random numbers. In my experience, that's also the only situation where you want such an API. For everything more complicated, you need special cases anyway.

Edit: FWIW, here is how I handle your specific example in my wrapper.

# file: random_numpy.np
def split(rng):
    return rng, rng

# file: random_jax.py
def split(key):
    return random.split(key)

# in implementation, rand is rng or key
rand, subrand = random.split(rand)
my_function(1.0, subrand)
rand, subrand = random.split(rand)
my_function(1.0, subrand)

Given lack of ecosystem agreement, we are not likely to forge a path forward for standardizing a PRNG API in the array API specification at this time. As such, I will go ahead and close this issue.

We can reopen/revisit if and when we have greater community consensus.

3. Copy-and-return-iterated pure functional (`rng, x = uniform(rng, 0, 1)`) (not implemented by anyone that I'm aware of)

This is such an insightful discussion. Thanks for having in in the open!

Just wanted to say that this is actually what we do at PyTensor (descendants of Theano).

Our functional API is next_rng, sample = uniform(rng). Then during compilation, if a user (or we automatically) state that when computing next_rng, rng can be mutated, we simply do what numpy does.

Otherwise, we copy the rng before using. This is wasteful of course, but useful for debugging, doing intermediate evaluations without incurring side-effects until we're ready to.

https://pytensor.readthedocs.io/en/latest/tutorial/prng.html

Navigation Menu

Repository navigation

Adding random number generation API #431

24 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Adding random number generation API #431

Description

Activity

leofang commented on May 18, 2022

jakirkham commented on May 18, 2022

rgommers commented on May 18, 2022

jakirkham commented on May 18, 2022

oleksandr-pavlyk commented on May 18, 2022

rgommers commented on Jun 1, 2022

rkern commented on Jun 1, 2022

rkern commented on Jun 1, 2022

rgommers commented on Jun 1, 2022

jakirkham commented on Jun 1, 2022

seberg commented on Jun 1, 2022

rgommers commented on Jun 2, 2022

rkern commented on Jun 2, 2022

shoyer commented on Jun 2, 2022

shoyer commented on Jun 2, 2022

24 remaining items

jakevdp commented on Jun 8, 2022

seberg commented on Jun 8, 2022

froystig commented on Jun 8, 2022

rkern commented on Jun 8, 2022

rkern commented on Jun 8, 2022

froystig commented on Jun 8, 2022

wangpengmit commented on Sep 3, 2022

rkern commented on Sep 3, 2022

ntessore commented on Mar 3, 2023

jakevdp commented on Mar 5, 2023

ntessore commented on Mar 5, 2023

kgryte commented on Jun 29, 2023

ricardoV94 commented on Jan 31, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions