Clean way to map randomness over a pytree

[greeneggsandyaml]

Hello authors,

Thank you for this wonderful library. I have a quick question: what is the cleanest way to map a random function over the leaves of a pytree? For example, let's say I want to add some N(0,1) Gaussian random noise to each parameter in a dict of parameters.

Concretely, suppose we have the following:

d = {'a': {'aa': jnp.zeros(4), 'ab': jnp.zeros(4)},
     'b': {'aa': jnp.zeros(5), 'ab': jnp.zeros(5)}}

I imagine that there is a way of doing something like:

key, *keys = random.split(key, 1 + len(jax.tree_leaves(d)))
d_new = jax.tree_map(lambda x, *keys: x + jax.random.normal(key=<the i-th key>, shape=x.shape, *keys), d)

or something similarly clean. Alternatively, if there's a better way of doing this, I'd love to know.

I figured out that it's possible to do

leaves, treedef = jax.tree_flatten(d)
new_leaves = []
for leaf in leaves:
    key, new_key = random.split(key)
    new_leaf = jax.random.normal(key=new_key, shape=leaf.shape)
    new_leaves.append(new_leaf)
new_d = jax.tree_unflatten(treedef, new_leaves)

and this works, but there must be a better way. I feel like I'm missing something obvious here, so I figured that I should ask.

Thanks for your help!

[jakevdp]

I imagine the best approach would be to use tree_structure, tree_unflatten, and tree_map:

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

key = random.PRNGKey(0)

d = {'a': {'aa': jnp.zeros(4), 'ab': jnp.zeros(4)},
     'b': {'aa': jnp.zeros(5), 'ab': jnp.zeros(5)}}

key, *subkeys = random.split(key, 1 + len(jax.tree_leaves(d)))
subkeys = jax.tree_unflatten(jax.tree_structure(d), subkeys)
result = jax.tree_map(lambda x, key: x + jax.random.normal(key, x.shape), d, subkeys)
print(result)
# {'a': {'aa': DeviceArray([-0.9407592 , -1.4780389 , -1.0911523 , -0.23124945], dtype=float32),
#        'ab': DeviceArray([ 0.6453762 , -1.1911646 ,  0.82499367, -0.35288557], dtype=float32)},
#  'b': {'aa': DeviceArray([ 0.8250121 ,  0.8801492 , -1.1567167 , -0.03485736, 0.41830325], dtype=float32),
#        'ab': DeviceArray([-2.46225   ,  1.622537  ,  1.429637  , -0.42747816, 1.1972864 ], dtype=float32)}}