[6gvONxR4sf7o]

> with jax.disable_jit(): ...

That's handy, and I hadn't seen it before, thanks.

It's been a bit, but I think the most frustrating errors were around mapping pytrees (like this issue https://github.com/google/jax/issues/9928). I'm not sure the exact solution, but the axis juggling and specifications were where I remember a lot of pain, and the docs (though extensive) were unclear. At times it feels like improvements are punted on in the hopes that xmap eventually fixes everything (and xmap has been in experimental for far longer than I expected).

Also the barriers where I couldn't disable jit. IIRC pmap automatically jits, so there was no way to avoid staging that part out. When it came to doing some complex jax.lax.ppermute, it felt more difficult than it needed to be to debug.

Next time I encounter something particularly opaque, I'll share on the github issue tracker.

[mattjjatgoogle]

Thanks for taking the time to explain these.

> It's been a bit, but I think the most frustrating errors were around mapping pytrees (like this issue https://github.com/google/jax/issues/9928).

We've improved some of these pytree error messages but it seems that vmap one is still not great. Thanks for the ping on it.

> Also the barriers where I couldn't disable jit. IIRC pmap automatically jits, so there was no way to avoid staging that part out.

That was indeed a longstanding issue in pmap's implementation. And since people came to expect jit to be "built in" to pmap, it wasn't easy to revise.

However, we recently (https://github.com/google/jax/pull/11854) made `jax.disable_jit()` work with pmap, in the sense that it makes pmap execute eagerly, so that you can print/pdb/etc to your heart's content. (The pmap successor, shard_map (https://jax.readthedocs.io/en/latest/jep/14273-shard-map.htm...), is eager by default. Also it has uniformly good error messages from the start!)

> Next time I encounter something particularly opaque, I'll share on the github issue tracker.

Thank you for the constructive feedback!

[6gvONxR4sf7o]

Thanks! One last thing, since I have your ear. The function transformation aspects of jax seem to make their way into downstream libraries like haiku, resulting in a lot of "magic" that can be difficult to examine and debug. Are there any utils you made to make jax's own transformations more transparent, which you think might be helpful to third party transformations?

Higher order functions are difficult in general, and it would be fantastic to have core patterns or tools for breaking them open.

[patrickkidger]

It sounds like you're concerned about how downstream libraries tend to wrap JAX transformations to handle their own thing? (E.g. `haiku.grad`.)

If so, then allow me to make my usual advert here for Equinox:

https://github.com/patrick-kidger/equinox

This actually works with JAX's native transformations. (There's no `equinox.vmap` for example.)

On higher-order functions more generally, Equinox offers a way to control these quite carefully, by making ubiquitous use of callables that are also pytrees. E.g. a neural network is both a callable in that it has a forward pass, and a pytree in that it records its parameters in its tree structure.

[6gvONxR4sf7o]

As a matter of fact, you’re preaching to the choir! Equinox is my go-to library for jax NN work!

[mattjjatgoogle]

You're right that downstream libraries have often tended to introduce magic (some more than others), and moreover one library's magic is typically incompatible with other libraries'. It's something that we're working on but we don't have much to show for it yet. Two avenues are:

1. as you say, exposing patterns and tools for library authors to implement transformations/higher-order primitives using JAX's machinery rather than requiring each library to introduce bespoke magic to do the same;

2. adding JAX core infrastructure which directly solves the common problems that libraries tend to solve independently (and with bespoke magic).

[6gvONxR4sf7o]

Thanks for the info! And I want to be clear that I make all these comments and questions out of a love for the tool. It's one of my favorite tools that I wish I could use at work. The design and community engagement are both fantastic.