JAX tree_* drops `None` elements

Somewhat confusingly (imho) flatten drops leaves that are None:

>>> jax.tree_leaves([1, None, 1])
[1, 1]

However flatten/unflatten on such a structure is symmetrical (leading me to believe JAX treats a None leaf as part of the structure):

>>> leaves, treedef = jax.tree_flatten([1, None, 1])
>>> jax.tree_unflatten(treedef, leaves)
[1, None, 1]

However multimap does not complain that the structures are different, and also skips the map fn if the left hand leaf is None:

>>> a = [None, 1]
>>> b = [2, None]
>>> jax.tree_multimap(lambda x, y: (x, y), a, b)
[None, (1, None)]

Is this desired behaviour? It doesn’t align with the other tree libraries (e.g. tf.nest and dm-tree) which treat None as a leaf, so it feels broken to me as a user of those libraries.

[skye]

Just my $.02, I would prefer if this were behavior were switchable, as I also find the current behavior surprising and hard to remember (I screwed up using pytrees with Nones not five minutes ago!) and I don’t like the vmap dance. But agreed this isn’t top priority.

[alvarosg]

+1 of documenting this more thoroughly in the doc strings of methods in tree_util. I don’t see any mentions to None in the tree_utils docs

And the notebook referenced earlier does not seem exist anymore.