changes coming to JAX's RNG – auto-parallelizable by default

[froystig]

Partitionable Threefry RNG upgrade

You can stop reading now if:

• you do not generate random numbers with JAX; or
• you already set the JAX config option jax_threefry_partitionable to True; or
• you only use the non-default unsafe_rbg JAX RNG algorithm.

What is happening?

• JAX's default RNG algorithm ("Threefry") is changing under the hood, to make random number generation efficiently auto-parallelizable ("partitionable"). This makes random numbers faster with multiple devices.

• The current behavior corresponds to setting the config field jax_threefry_partitionable to False, its current default value. The new behavior corresponds to setting jax_threefry_partitionable to True, its future default value.

• Some time soon, we will set out to change the default value of the configuration flag jax_threefry_partitionable from False to True. This will cause a one-time change in the random values generated from a given RNG key. See code below.

• Want to ensure this goes smoothly? Try flipping jax_threefry_partitionable to True today to detect any issues in your code ahead of the upgrade. See code below.

What do we mean by a one-time change to random values? Here is today's default behavior:

$ JAX_THREEFRY_PARTITIONABLE=False python -c '
> import jax
> print(jax.random.randint(jax.random.key(72), (), 0, 10))'
3

And here is what will happen when the default setting changes soon:

$ JAX_THREEFRY_PARTITIONABLE=True python -c '
> import jax
> print(jax.random.randint(jax.random.key(72), (), 0, 10))'
7

Same key, different generated value when JAX_THREEFRY_PARTITIONABLE=False versus JAX_THREEFRY_PARTITIONABLE=True.

JAX's RNG will remain deterministic. We try to rarely change the output of JAX's pseudorandom functions. That said, our API policy promises stability in distribution, not in value. This particular change is broad-based, and is the first of its kind in a long time, so we're drawing extra attention to it.

This change will break code that depends on specific RNG keys generating specific RNG values. Common examples include reference tests, high-variance randomized tests, or a machine learning experiment that depends on random values (e.g. model initialization) that you want to reproduce exactly.

Non-partitionable Threefry will later be deprecated. That is, at some point after we've upgraded the default value of the jax_threefry_partitionable setting, we will deprecate the flag entirely.

Who is affected?

JAX supports several RNG schemes. Its current three built-in modes are called threefry2x32, rbg and unsafe_rbg. These modes can be set using the impl argument to jax.random.PRNGKey and jax.random.key, or with the configuration flag jax_default_prng_impl.

You are affected if you use either threefry2x32 or rbg. Specifically:

• The random values generated from a key will change for you if you use threefry2x32.
• Key derivation will change in rbg mode. That is, the keys generated from a specific key, using jax.random.split or jax.random.fold_in, will be different than before. In turn, random values generated from such derived keys will change.

Opting in early

To try the upgrade now, you can set the configuration flag jax_threefry_partitionable to True in your code explicitly. This can be done with the environment variable JAX_THREEFRY_PARTITIONABLE=True, the command-line flag --jax_threefry_partitionable=True, or programmatically, using jax.config.update or the jax.threefry_partitionable context manager. For example:

jax.config.update("jax_threefry_partitionable", True)

# OR

with jax.threefry_partitionable(True):
  ...

[stergiosba]

There is a typo in what is happening:
Change jax_threefry_patitionable to jax_threefry_partitionable if you want.

Thanks for the heads up!

[jakevdp]

Done - thanks!

[jakevdp]

This is unrelated to the topic here - please open a new discussion for this question.