An advice on missing pytree functionality

[hr0nix]

There are several pytree operations that, IMO, would be very useful if present in jax.tree_util:

• Pytree merging
  Say I have 2 pytrees with different structures: tree1={"a": 2, "b": {"c": 1}} and tree2={"a": 3, "d": 5}. I would like to do a tree_map on these trees with some function f so that the map result is {"a": f(2, 3), "b": {"c": f(1, None)}, "d": f(None, 5)}, i.e. we pass None to f for leafs that are missing in one of the trees. This would, for instance, simplify model surgery, allowing to combine parameters from multiple models.
• Some form of pytree filtering, i.e. building a subtree of a pytree based on some provided node selector function. It would greatly simplify things such as metric computation, e.g. when I need to compute norms of some of the parameters or gradients of my model and return them as a tree while preserving the structure.

Are there any existing libraries built on top of jax.tree_util that can do something similar? Or, perhaps, there is a reasonably easy way to build something like that from existing primitives (I can't think of any)?

[patrick-kidger]

• On PyTree merging: note that dictionaries with different keys are completely different PyTrees. So asking to merge {"a": *, "b": *} and {"a": *, "d": *} isn't defined, PyTree-wise. (If you want, you can of course do something else, that is dictionary-specific, instead.)

  That aside and for completeness, you might like equinox.combine, which does allow you to merge general PyTrees, by filling in Nones.

• On PyTree filtering: see equinox.filter, which does exactly this.

• On libraries that add extra tree manipulation routines: see this page of the Equinox docs, which includes a lot of useful functions.