Comment by 6gvONxR4sf7o

That's true, and is a massive part of what I love about JAX, but they also form barriers in weird parts of your code, preventing standard introspection tools, which is the single thing I hate about JAX. The errors are amazingly opaque.