Trace detection when args are pytrees

[pierreglaser]

Hey all!

Currently, find_top_trace (see
https://github.com/google/jax/blob/d4965cd33549ebe66cd474617366aa59b894d676/jax/_src/core.py#L1285-L1296
)

seems to assume that xs is a tuple of (Tracers of) jax Arrays - my understanding is that the detection mechanism just breaks down if xs is, say, a tuple of pytrees which themselves contain jax tracers. This phenomenon surfaces when defining a custom batching primitive for a function taking a pytree as input. Curious if any core dev has any thoughts on this? Does that make sense? Has support for trace detetion from pytrees ever been considered? I should be able to provide a reproducer if necessary.

[jakevdp]

Hi - you're correct that this function assumes the input is a flattened list of array-like objects. At the level of the stack where this is called, pytrees have already been flattened away, so there shouldn't be any other structure here. This is an internal utility that you shouldn't need to call directly – can you say more about what code you were running when you ran into this problem?

[pierreglaser]

so there shouldn't be any other structure here

I assume this holds because find_top_trace is called when evaluating "atomic" primitives like sin_p and others, which (unless I'm missing something), always take arrays as input (and not pytrees)

But my use-case is different: essentially, I wanted to write a small compatibility bridge between "pytorch functions" (e.g. functions manipulating torch tensors) and jax (I was using such functions with other libraries that expected jax functions and not pytorch functions). Thus, I created a small decorator utility that transforms a function expecting torch tensors into a one expecting jax arrays. Under the hood, the function outputted by the decorator just converts the jax arrays into torch tensors using dlpack and then passes them to the original pytorch function. Since I also wanted the function outputted by the decorator to be compatible with jax.vmap, the decorator also registers a custom batching rule for this function, which similarily perfoms jax->torch tensor conversion and then calls torch.vmap under the hood with the correct in/out axes.

This all works actually quite well, and even supports nested jax.vmap calls. However, when the original pytorch function takes as input, say, a dict of torch tensors (meaning the decorated function will expect a pytree of tensors), then batching traces are incorrectly flagged as evaluation traces by jax itself when evaluating vmap-ed version of the function outputted by the decorator, and the reason (IIUC) is the behavior of find_top_trace that I described above: in that context, there will be no BatchTracer in the xs, but a pytree of BatchTracer. I never call find_top_trace directly :-)

[pierreglaser]

Mmmh. And then unflatten the input back before feeding them to the torch function? That would mean pass around the treedef?

[jakevdp]

Yes, in general this is how JAX internally handles pytree arguments. All internal code is built to handle flat sequences of buffers, and all interfaces to the outside world flatten pytrees on the way in, and unflatten on the way out. It means that most of the JAX internals don't need to know anything about pytrees, which makes implementations much simpler!

[pierreglaser]

I see. Thank you for taking the time to answer, that was super useful!

[patrick-kidger]

You may also like Equinox's filtered primitives, which do offer a pytree-based interface: https://github.com/patrick-kidger/equinox/blob/main/equinox/internal/_primitive.py

Although this is private/undocumented (just like JAX's primitives), so y'know use at your own risk.

[pierreglaser]

Thanks @patrick-kidger for chiming in!