Conversions between JAX PRNGKey, integer seed and NumPy RandomState/Generator

[ethanluoyc]

Hi,

Is there a recommended approach for conversions back and forth JAX PRNGKey and other random number generators?

Frequently, the user needs to use random number generators in different libraries (NumPy, or other libraries where they expect an integer seed). Imagine the case that I have a function that takes a JAX PRNGKey. Inside the function, I would need to split the JAX PRNGKey and use it to seed other libraries (with NumPy RandomState, integer seeds). However, since the other libraries may not support JAX PRNGKey, I would need to do some conversions.

For example, here is my approach to converting PRNGKeys to integer seeds for seeding OpenAI's Gym environment. I wrote

def fun(seed):
  random_key = jax.random.PRNGKey(seed)
  agent_key, env_key = jax.random.split(random_key)
  env = gym.make()
  seed = int(jax.random.randint(env_key, (), 0, jnp.iinfo(jnp.int32).max))
  env.seed(seed)
  # make some agent
  agent = make_agent(agent_key)

as opposed to

def fun(seed):
  agent_key = jax.random.PRNGKey(seed)
  env = gym.make()
  env.seed(seed + 1)
  # make some agent
  agent = make_agent(agent_key)

Since I heard the later approach may suffer some issues with correlations in the seeds, but I am not sure if this would be true for JAX.

Some libraries expect NumPy RandomState/Generator to be used, in those cases I suppose it's safe to simply use the PRNGKey to seed the RandomState, so I suppose it's ok to just do

key = jax.random.PRNGKey(seed)
rs = np.random.RandomState(key)

I would love to know the approach I am using above is sound, or is there an even safer approach to do these conversions. Thanks!

[jakevdp]

I don't know of any definitive answer regarding this.

For numpy in particular, the new np.random.default_rng object has the ability to seed from an arbitrary array of integers, so I would think the best option would be to pass the jax seed directly as an array:

from jax import random
import numpy as np
key = random.PRNGKey(0)
rng = np.random.default_rng(np.asarray(key))

[ethanluoyc]

Yeah, this is a really neat thing to do! I guess this would the recommended approach for interop with NumPy's Generator.

Some of the libraries I use are still only supporting np.random.RandomState,
I just found out that a similar thing can be done for RandomState, the following works

import jax.numpy as jnp
import jax
import numpy as onp

key = jax.random.PRNGKey(0)
onp.random.RandomState(key) # >> RandomState(MT19937) at 0x7F5202AE5440

[aswanthkrishna]

can i use this inside a jitted function?

[ethanluoyc]

I don't think so. The creation of onp.random.RandomState should be outside a jitted function, otherwise I think it would not work.