Adding tree_util.stack_leaves() and tree_util.unstack_leaves()

[ayaka14732]

• stack_leaves: Stack the leaves of one or more PyTrees along a new axis.
• unstack_leaves: Unstack the leaves of a PyTree.

References:

• https://docs.liesel-project.org/en/v0.1.4/_modules/liesel/goose/pytree.html#stack_leaves
• https://gist.github.com/willwhitney/dd89cac6a5b771ccff18b06b33372c75?permalink_comment_id=4634557#gistcomment-4634557
• https://github.com/ayaka14732/llama-2-jax/blob/ab33e1f15489daa8b9040389c77e486cd450e461/lib/tree_utils/__init__.py

[jakevdp]

To be clear, are these the semantics you have in mind?

def stack_leaves(pytrees, axis):
  return jax.tree.map(lambda *xs: jnp.stack(xs, axis), pytrees)

[ayaka14732]

To be clear, are these the semantics you have in mind?

Yes

[jakevdp]

For something like this, I'd probably lean toward recommending users implement what they need via existing API composability, rather than providing a new API for something that can already be pretty succinctly expressed. What do you think?

[ASEM000]

Maybe adding tree util cookbook would be useful? @jakevdp

[jakevdp]

A pytree cookbook would be an interesting idea! This idea also came up in #20594. @ayaka14732, is that something you'd be interested in thinking about?