How to unflatten a PyTree one level?

[adonath]

I have some code to wrap selected nodes in any given PyTree. The way I achieved that was to traverse the tree by flattening one level, wrapping the node and unflatten again, with the wrapped node. For the implementation I relied on the private _registry_with_keypaths, which gave me access to the node specific flatten and unflatten methods. Here is the corresponding code:

https://github.com/adonath/clouseau/blob/main/clouseau/jax_utils.py#L59

In JAX 0.6 the _registry_with_keypaths registry has been removed and I am wondering how to achieve the unflattening of the node from one level using the public API. While in principle I can just instantiate the node from the children and aux data, there is no information on how the aux data and children are mapped in the __init__ method. The behavior is hard-coded in the registered unflatten function.

So the concrete questions are:

• How to unflatten a PyTree one level?
• Is there some way of retrieving the flatten/unflatten methods from the public default_registry?

I would be happy for any hints or suggestions for alternative solutions!

[jakevdp]

You can do this:

jax.tree.flatten(x, is_leaf = lambda node: node is not x)

[adonath]

Thanks @jakevdp! I can confirm the following pattern works as intended:

children, treedef = jax.tree.flatten(x, is_leaf = lambda node: node is not x)
children = ... # modify children
node = treedef.unflatten(children)