Converting Pytree params to a vector

[dptam]

Hi, I was wondering how to convert a pytree of params to form the a single vector with all the parameters flattened, and then mapping this vector back to the pytree of params.

1. At the moment, I flatten the pytree of params, and then concatenate the list of values. However, to map this single vector back to the pytree of params, I have to first jnp.split this vector into the a list of individual parameter shapes, reshape each parameter to its original size, and then unflatten this list to form a pytree. It seems jnp.split might not be jittable? I was wondering if there was a more efficient or jittable way to do this operation?
2. Also, it seems from jax.tree_utils do not keep dict key order #4085 that jax.tree_map does not necessarily keep the order of keys in the dictionary. This causes the order of parameters to be different if I first map the Pytree of params to do some operation, and then flatten it. Is there a way around this at the moment?

Thanks for your help.

[PhilipVinc]

jax.flatten_util.tree_ravel does exactly this. It returns the flattened vector + a function to unravel (or unflatten) the vector.

[dptam]

Thanks for your help. This is exactly what I'm looking for.

[ddrous]

@PhilipVinc This function doesn't appear to exist in JAX anymore. It is super valuable. Is there a replacement ?

[PhilipVinc]

It's called ravel_pytrre those days