How does jit compute cache keys and hits for generic PyTrees?

[KeAWang]

From the docs it seems that jax.jit computes the cache keys for static args differently than the dynamic args.
Suppose we have a function

def f(x, y):
    ...

Let's say we jit f such that the resulting function expects x to be a pytree and y to be static (which may or may not be a PyTree).

My understanding is that cache hits for y are computed using its __eq__ and __hash__, and that this is true regardless if y is a PyTree, a pure user-defined function, or some other hashable object.

However, what does it use to compute cache hits when x is a generic PyTree? I know that if x is a jax.Array, it's based on its dtype and shape. But what if x is a tree_util.Partial or a pair of Booleans? Or some other arbitrary PyTree that contains arbitrary leaves and auxilliary data?

[jakevdp]

The rough mental model is this: if you pass a pytree y to the function, then internally the JIT machinery computes

children, treedef = tree_util.tree_flatten(y)

The children are treated as dynamic arguments, while treedef (and all the pytree auxilary data it contains) is treated as a static argument, with its hash being rolled into the JIT cache key.

[KeAWang]

Gotcha. So even though technically PyTree leaves can be any object type (e.g. a class, a callable), jax.jit requires dynamic PyTree arguments to have traceable leaves.

[jakevdp]

Exactly – I've added some documentation of this in #15007