Understanding vmap + make_rng

[JamesAllingham]

I tried to use jax.vmap inside a linen.module where I was calling self.make_rng and I found the interaction surprising. Here is a minimal example:

class MyModule(nn.Module):
    @nn.compact
    def __call__(self):
        def vmapped_fn(i):
            return jax.random.normal(self.make_rng("my_key"), x.shape)

        return jax.vmap(vmapped_fn)(jnp.arange(3))

module = MyModule()
result = module.apply({"params": {}}, rngs={"my_key": jax.random.PRNGKey(0)})

print(result)
>>> [[-1.2807428]
>>> [-1.2807428]
>>> [-1.2807428]]

I was initially surprised to see that each call to jax.random.normal receives the same key, despite the fact that self.make_rng is supposed to return a different key each time it is called. I guess that because self.make_rng is stateful, all bets are off for how it interacts with vmap? Or is there something else going on here?

In this case, the fix is simple, just fold in an index:

class MyModule(nn.Module):
    @nn.compact
    def __call__(self):
        def vmapped_fn(i):
            return jax.random.normal(jax.random.fold_in(self.make_rng("my_key"), i), x.shape)

        return jax.vmap(vmapped_fn)(jnp.arange(3))

module = MyModule()
result = module.apply({"params": {}}, rngs={"my_key": jax.random.PRNGKey(0)})

# Print the result
print(result)
>>> [[-0.28846943]
>>> [-1.3579164 ]
>>> [ 0.62705433]]

Is this the best fix, or should I use linen.vmap somehow?

[chiamp]

With nn.vmap, you could do something like this:

class RandomModule(nn.Module):
  @nn.compact
  def __call__(self, x):
    return jax.random.normal(self.make_rng("my_key"), x.shape)
class VmapRandomModule(nn.Module):
  @nn.compact
  def __call__(self, x):
    batch_random_module = nn.vmap(
        RandomModule,
        split_rngs={'my_key': True})
    return batch_random_module()(x)
x = jnp.arange(3)
a = VmapRandomModule()
a.apply({}, x, rngs={'my_key': jax.random.PRNGKey(0)})
>>> Array([1.6090478 , 0.16792756, 0.1328707 ], dtype=float32)

Also I think you had a typo in your code, since the variable x isn't defined:

class MyModule(nn.Module):
    @nn.compact
    def __call__(self):
        def vmapped_fn(i):
            return jax.random.normal(jax.random.fold_in(self.make_rng("my_key"), i), i.shape) # instead of x.shape

        return jax.vmap(vmapped_fn)(jnp.arange(3))

In this case, it would fold in the same rng key and you'd only get a different output if the input i is different. In your example it works because you use jnp.arange(3), but if you used something like jnp.ones(3), it will output the same value for all 3 numbers.

[JamesAllingham]

Thanks @chiamp -- that all makes sense!