New DeprecationWarning unclear for jax.tree.map

[liukidar]

Hello everyone :)

The following code snippets rises a DeprecationWarning in the last version of JAX which, to my understanding, shouldn't be there. The code is

tree = jtu.tree_map(
            lambda _, y: y,
            filter,
            tree,
            is_leaf=lambda x: isinstance(x, Node),
        )

Basically, I filter tree by selecting only the subtrees with selected Nodes as a root, based on filter. Here, the idea is that filter has been modified by replacing all the excluded Nodes with None.

In particular the warning I get is

DeprecationWarning: In a future release of JAX, flatten-up-to will no longer consider None to be a tree-prefix of non-None values, got: TypeA.

To preserve the current behavior, you can usually write:
  jax.tree.map(lambda x, y: None if x is None else f(x, y), a, b, is_leaf=lambda x: x is None)

which makes me wonder if the code I wrote is behaving as expected, since it relies on the fact the None values are excluded from consideration, while the message seems to indicate that now they are included (and will only be excluded in future versions)?

To remove the warning I can write:

module = jtu.tree_map(
            lambda x, y: None if x is None else y,
            self.filter.get(),
            module,
            is_leaf=lambda x: x is None or isinstance(x, BaseParam),
        )

which to me should be functionally identical to the original code.

Am I missing something?

[jakevdp]

The warning arises because None is currently being treated as a leaf in the context of tree_map, and as an empty tree in other tree_util context. Your code is relying on None implicitly being treated as a leaf, and in the future your code will break unless you explicitly handle that. There's a bit more discussion of this in the JAX 0.4.29 Change Log.

To be clear, your new code is functionally identical to the old code in current JAX releases, but will not be functionally identical in future releases, because the semantics of your original code will change (thus the deprecation warning).

Does that make sense?

[liukidar]

I think in my code I am relying on Nones NOT being treated as leaves (i.e., map_fn is not called on them), since I want to remove the subtrees of module for which the corresponding root element in filter is None. How would the behaviour of the first code snippet above be different in newer versions? (From my understanding both should ignore None values).

[jakevdp]

Could you share a full example of what code you're executing, including defining the inputs to the function? It's easier to discuss the behavior if the example is more concrete.