Using keys to extract Pytree values

[SNMS95]

Hi guys,

I am trying to normalize the leaves in one tree based on the values of the leaves in a similar tree. The problem is that I have to do this based on a condition. The code is somethign as follow:

# Create reference tree
tree1 = {'layer1':{'w': np.array([1., 2., 3.])},
         'layer2':{'w': np.array([2., 4., 6.])}} 
flat_tree1, unflatten_op = jax.tree_util.tree_flatten(tree1)
# Create a random tree
vector = jax.random.normal(jax.random.PRNGKey(0), flat_tree1.shape) 
tree2 = unflatten_op(vector) # random tree

flatten, tree_def = jax.tree_util.tree_flatten_with_path(tree1)

for key, value in flatten:
    if 'some_string' in jax.tree_util.keystr(key):
        tree2[key] = tree2[key] / np.linalg.norm(value)  # error
    else:
        tree2[key] = tree2[key] * np.linalg.norm(value, axis=0)  # error

But it seems I can't access the the values of pytrees with the key?
Is there any way to do it?

[chrisflesher]

Can you just use getattr(tree2, key)?

Also there is the equinox package that I've found useful for manipulating pytrees. It has a tree_at function that does what you want.