Getting all values in a pytree

Hi, sorry if this is a question with an obvious answer, but how do I get all the values in a pytree? I tried using tree_flatten, but it gives me the device arrays, not the values they contain (For extra context I want to implement regularizers using jax and need to sum all the squares/absolute values of the neural network parameters)

[jakevdp]

It does still exist:

from jax.flatten_util import ravel_pytree

It looks like it’s not listed in JAX’s documentation

[mattjj]

It works to ravel things, but that might actually be less efficient because it probably involves concatenating all the values into a single contiguous buffer. It might be better to do something like

def l2_normsq(x):
  leaves, _ = tree_util.tree_flatten(x)
  return sum([np.sum(leaf ** 2) for leaf in leaves])

(where sum is just Python’s builtin.)