The main idea (both pro and con) with Jax is its "automatic performance" model. Essentially, Jax uses a tracer down to XLA which then does some linear algebra optimizations to generate code. For the DSL case where Jax is generally used, this tends to be "good enough" for a good number of use cases. The issue is that when it's not "good enough", it doesn't give package developers much recourse to really dig into the details and optimize for the things which the compiler (and compilers in general) is unable to do.

For example, we released the library SimpleChains.jl that was optimized for a lot of cases for small neural networks (https://julialang.org/blog/2022/04/simple-chains/). There was discussion about Jax, and we saw that for any case where individuals had a modern multi-core CPU threaded machine setup, SimpleChains.jl had about a 10x performance improvement on CPU over the Jax equinox library. This shouldn't be surprising to anyone because the details of how that was achieved was quite clear (memory handling, manual SIMD bits, etc.), though those of course are not the kinds of things you get "automatically" with either language. The main difference of having a full programming language instead of a DSL is that such code can be written to manually optimize what the compilers are not doing.

And this is the pattern that tends to repeat. For small static-like differential equation solves, Jax is fine, but for example the PDE code optimizations done in this tutorial (https://diffeq.sciml.ai/dev/tutorials/faster_ode_example/#Ex...) are completely incompatible with the Jax programming model and accelerate the PDE solve by more than 10x. And if you look into detail at things showing Jax is fine for PDEs, they are using the slower code style and comparing to that slower Julia code, not the faster iteration form. Of course, these are optimizations most people don't know to do so they are the kind of extra order of magnitude you see in Julia packages over Jax, but not as much in user code.

Another repeat of this pattern is the sub-optimality of vmap. It doesn't run code completely independently so it's not optimal in CPU cases where you'd want to effectively use multithreading or MPI, but it's also not generating kernels so it's not optimal in the case of GPUs (for example, we have a new GPU-based ODE solver coming out soon that outperforms the older Julia form and the Jax form by like 100x. Again, no surprises here because special-built GPU kernels have already demonstrated this as possible from C++ directly in this context). So vmap will always get you parallelism, but when you build a library you have to leave it behind at some point if you really want to keep optimizing.

I think what it does is admirable. If you can get a code to compile with Jax (which is a big if, it does require pure functions and a lot of extra things Python code normally doesn't satisfy) then it does get you something relatively performant. But there's a reason why what's mentioned here is Flux and Zygote, not SimpleChains and Enzyme, the real tools the Julia community is pushing towards with the full mutation support and the full speed behind it. Jax can compare in the former case, but not the latter case.