Jax vs. Julia (Vs PyTorch)

To contribute to an existing repository, you can ]dev the package to download the git repository, @edit a function to open the source in an editor, and push the changes to a PR in about 10 minutes. This makes it much easier for package users to turn into package developers.

I've contributed to a bunch of Julia libraries that if I still used python, I wouldn't have contributed to the equivalent python library because contributing to Julia packages is at least an order of magnitude easier.

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.

No. I suggested I myself have seen gatekeeping and referenced your comment as example. Make of that what you will, but don't extrapolate it to a conspiracy theory I didn't suggest.