It seems like the ecosystem is still dominated by PyTorch, is Jax supposed to be a competitor? Any signs of Jax taking over PyTorch anytime soon? Is it perhaps too early for its time? Or is there a critical flaw in the underlying design?
The beauty of JAX is that basic usage is basically a single function: `grad`.
You just write whatever Python function you want and can get the derivative/gradient of it trivially. It gets a bit trickier when you need more sophisticated numeric tools like numpy/scipy, but in those cases it's just about swapping out with a JAX version of those.
In this sense JAX is the spiritual success to Autograd. However the really amazing thing about JAX is that not only do you get the autodiff for basically free, you also get very good performance, and basically GPU parallelism without needing to think about it at all.
PyTorch is an awesome library, but largely focus on building Neural Networks specifically. JAX should be thought of a tool that basically any Python programmer can just throw in there whenever they come across a problem that benefits from having differentiable code (which is a lot of cases once you start thinking about differentiation as a first class feature).
I don't get the point of this distinction - JAX was developed specifically for ML. What else is it being used for right now?
It’s probably also useful in population and metaheuristic scenarios where the optimisation objective can be described mathematically, allowing you to make use of GPGPUs, and if possible first and second order derivatives.
Although, I've mainly heard about Julia in that context, not Jax.
But in general, I would suspect youth.