Randomness within a vmap?

[khdlr]

Hi everybody,

this is probably a stupid question, but I just can't seem to get it right. I want to use random sampling within a function that I will vmap later on. The odd thing is, that after vmapping, the random samples will be constant over the mapped axis. How can I configure the vmap in such a way that this is not the case?

Here's a MWE of what I mean:

key = jax.random.PRNGKey(42)

def add_gaussian(inp, key):
    return inp + jax.random.normal(key, inp.shape)

add_gaussian(jnp.zeros([]), key)
# Returns -0.18471177

vectorized_add_gaussian = jax.vmap(add_gaussian, [0, None], 0)
vectorized_add_gaussian(jnp.zeros(5), key)
# Returns [ -0.18471177, -0.18471177, -0.18471177, -0.18471177, -0.18471177 ]

I would like the last function call to output 5 different values. One solution would obviously be to split the rng-key and vmap over that as well, e.g.

vectorized_add_gaussian = jax.vmap(add_gaussian, [0, 0], 0)
subkeys = key.split(5)
vectorized_add_gaussian(jnp.zeros(5), subkeys)
# Returns [ 0.48309308, -1.1981592 , -0.7775587 ,  1.2407155 , -2.313012 ]

Is this the recommended way to handle the situation? Or is there a more elegant way of doing it? Manual vectorization would be possible without splitting the PRNG key, so I'm wondering whether vmap can do the same?

[hawkinsp]

It's a good question!

The key thing to understand about JAX PRNG keys is that they are just values, and they act like all other values with respect to transformations. So, if you want to have each element of a vmap have a different PRNG key, you should split the key and vmap over the split key. Indeed, that's exactly what you should expect from the semantics of vmap: anything else would be surprising.

This does require you to be explicit about threading PRNG state, but that's a good thing! If you did want to use the same PRNG key for each element of a vmap, it's clear how to do that (don't split), and if you want to use different PRNG keys for each vmap element, it's clear how to do that also (split).

Aside: We've long wanted to make it harder to accidentally reuse PRNG keys, so in the future we might require that you do something explicit to say "yes, I really want to use this PRNG key multiple times" in this sort of situation.

Hope that helps!

[khdlr]

Thanks for the quick and extensive reply! Thinking about the PRNG key as "just another value" indeed clarifies this.

If I may ask a follow-up question: Are there any performance-implications to this? I.e. is the "manually vectorized" code below faster than "auto-vectorized"? Or does calling jax.random.normal(key, [5]) perform similar RNG-state-splitting operations internally anyways?

# Manually vectorized
jax.random.normal(key, [5])

# Auto-vectorized
subkeys = jax.random.split(key, 5)
jax.vmap(jax.random.normal)(subkeys)

[Logon27]

I know this is an older question, but I searched around a bunch for a code example of the question asked here. I never really found an actual code example. So I figured I would contribute a simple example below to illustrate how this could be done. Hope this helps any future readers.

import jax
import jax.numpy as jnp
from jax import random

# Generate a starting key
key = jax.random.PRNGKey(0)
# Just an array of constant values as an example.
my_arr = jnp.array([1, 1, 1, 1, 1])
# Generate the same number of keys as the array size. In this case, 5.
subkeys = random.split(key, my_arr.shape[0])
# Printing subkeys and array just for visual purposes
print(f"subkeys: {subkeys}\n")
print(f"my_arr: {my_arr}\n")
# Vmap the list of subkeys to add a random number to each element of my_arr.
vmapped_result = jax.vmap(lambda x, y: (jax.random.normal(x) + y), in_axes=(0,0), out_axes=0)(subkeys, my_arr)
# You could replace the above lambda with your own custom function that takes in the subkeys and some array to manipulate randomly.
print(f"vmapped random values added to my_arr: {vmapped_result}")
# Prints...
# vmapped random values added to my_arr: [ 0.07951337  0.261127    2.6133053  -0.5701252   0.44482076]