How to access the default is_leaf function

[justindomke]

Many functions in jax.tree_util allow a user-specified is_leaf function. I'd like to override the default behavior in only a few cases. So it would be convenient to call the default is_leaf function. But how can I do that?

[jakevdp]

No, there is no default is_leaf function that you can override, rather the is_leaf function you provide overrides the normal tree-flattening logic.

If you're curious about how pytree flattening is implemented, you can find the implementation here: https://github.com/tensorflow/tensorflow/blob/949b63fa788411e4ddba84bbd70385aedd4ed4a2/tensorflow/compiler/xla/python/pytree.cc#L186 (Here leaf_predicate is the user-defined is_leaf function, if provided).

Roughly, the logic here is:

• if the user provided an is_leaf function, check if it returns True, and if so, return the node.
• otherwise, if the node's class is registered as a pytree, call into its flattening method
• otherwise, if the node is a hard-coded container type (tuple, list, dict, etc.) flatten each of its contents
• otherwise, the node is a leaf, and return it.

[justindomke]

I understand that. But in my function, I want to call the default is_leaf function.

If you want to understand the reason, I just want the functions to perform as normal, except treating None as a leaf. So I would like to create this function:

def is_leaf(x):
    if x is None:
        return True
    return default_is_leaf(x)

My problem is that I don't know how to call default_is_leaf. That's what I'm asking how to do.

[francois-rozet]

If x is a leaf, the value returned by is_leaf does not matter. So your is_leaf function should return False when you want the default behavior.

def is_leaf(x):
    return x is None

[jakevdp]

(sorry, I missed this originally because it was posted as a second answer, rather than a threaded response to my initial answer. @francois-rozet is correct)

[francois-rozet]

Thanks @jakevdp. By the way, if you absolutely need a function that returns True for "default" leaves, you can "hack" the jax.tree_util.all_leaves function:

def is_default_leaf(x):
    return jax.tree_util.all_leaves((x,))

This is pretty fast as it does not depend on the "depth" of x.