jax._src.flatten_util - raveling and unraveling.

[abhinavprao]

Hi I am trying to ravel a parameter - pytree into a single 1D array and rebuild it. Found the flatten_util module helpful for the raveling part. But to unravel we obviously need the tree_def and unravel_list (which I am assuming is being used to split and reshape into leaves). How do I get these from the HashablePartial object produced by the ravel_pytree function? Also is there a clean way to vmap both of these for batches?

[jakevdp]

jax._src.flatten_util is a private module and you should avoid using it directly, as it's not guaranteed to be stable between releases (See https://jax.readthedocs.io/en/latest/api_compatibility.html for information).

Can you give more detail on what precisely you're trying to achieve? There may be public APIs to do what you have in mind. I'm having trouble understanding what "ravel a parameter - pytree into a single 1D array and rebuild it" means.

[abhinavprao]

Sorry about the lack of clarity. I realized flatten_utils is not complete and closed my own thread. Here is what I am trying to achieve:

1. I am trying to train a hyper-network, which takes the 'parameters of a main-network' as input. Thus, I need to unravel all the parameters of my main-network into a single 1-D array (which my hyper-network can accept as input)

2. Then the hypernetwork outputs new parameters for my main-network. Thus, I need to unravel this output array back into a pytree with same structure as the original main-network parameters.

3. Parts 1 and 2 are easy to achieve with functions already in flatten_utils and I have finished this part.

4. The part I need help with is to be able to do this in a batch. Consider my original input pytree (parameters of main-network) are batched, i.e. each of my leaf's first dimension is batch_size. Now when raveling the pytree instaed of a 1D array I need a (batch_size, n_params) sized array, which can be unraveled back into a pytree.

Sorry for the long description. Do let me know if there are public APIs that do this. Thanks!

[jakevdp]

There's no public API for exactly this, but I would recommend writing a function to do this based on tree_flatten and concatenate; then the batched version would be just a simple vmap of the original function.

[jakevdp]

Actually, I forgot that jax.flatten_util.ravel_pytree is a public function. I was thrown off by the _src in your question, because everything in a _src path should be considered a private API.