I find JAX really exciting. The idea of numpy with autograd is exactly what Pythonistas want. The elephant in the room, though, is “why not Pytorch?”
Everyone knows JAX is what Google realised Tensorflow should have been when they realised how much of a joy Pytorch was to use. I actually think JAX does offer some advantages, not least true numpy interoperability. However, not mentioning *torch a single time in the blog post seems a little disingenuous for a Google-owned deep learning enterprise.
What is more some of the deep learning libraries on top of JAX (like flax, linen) are close to a one-to-one copy of PyTorch. However the underlying implementation is in some ways more interesting than the one of PyTorch and naturally accomodates higher derivatives and also the way that optimization and state initialization is implemented is more functional and "theoretically sound".
For me the main point is that JAX is slightly lower level than Pytorch and has nicer abstraction (no need to worry about tensors that might not store a gradient or wetehr you are on the GPU: it eliminates lots of newcomers bug) which makes it a great fit to build Deep learning frameworks, but also simulations, on top of it.
JAX and Pytorch have somewhat different scopes though... JAX is concerned about much more than the kinds of neural nets we write today. It's a general system for expressing and transforming numerical programs, and the devs are as genuinely excited about e.g. scientific programming, probabilistic modeling, etc. as they are about NNs.
A technical reason for "why not pytorch" is that JAX was also built in part to expose and leverage the power of the XLA compiler, which is at least for the moment a pretty uniquely powerful tool for producing efficient, highly-scalable accelerator code.
I should underline that this is a friendly community of peers though: there is a lot of respect for Pytorch, which in turn was certainly influenced by the original Autograd that many of the JAX devs also worked on. JAX (and its fancier sibling Dex) beyond being useful tools are also still research projects in and of themselves seeking to advance our ideas on how to write expressive, powerful numerical code on modern architectures.
Pytorch might be suboptimal but it's what's already there e.g. for Pyro. (PyMC3 uses Theano; can Theano drive JAX?)
OTOH I'm not sure most people know what tasks are GPU-worthy or not. I haven't the slightest idea of why MCMC/Variational Bayes is amenable to GPU speedups and Persistent Homology isn't.
JAX enables using (parts of) existing numpy codebases in disciplines other than deep learning. Autodiff and compilation to GPUs are very useful for all kinds of algorithms and processing pipelines.
No one does. (Also, that's a flax story, not a jax story. A bit confusing, but think of jax like numpy as flax like pytorch -- jax is the low level, flax is the high level.)
I think it's important to have global scope names. "biggan.discriminator.3.conv1.kernel.b" is perfectly sensible: it's the bias value of the kernel for the convolution of the third block of your biggan discriminator.
Everyone tries to treat model variables as interchangeable nameless parts. I hate it. Every variable has a global name, conceptually.
A global name also has other benefits. It becomes far easier to create EMA weights, for example, since you can filter the variables by name.
No really, it force you into pure, functional patterns. Frameworks build on top of JAX have done a lot of experiments to find a good way to deal with states while keeping the code clean and not too magicy.
Their is no clear best pattern at the moment. My favorite would be Flax and their recent Linen API which is a refined effort that pays off when using their framework.
The description says it's autograd + XLA. So I assumed it always compiles to GPUs via XLA.
But, had a look in the code and jax has cublas and RoCm blas, and it looks like there is a flow where it uses the gpu directly, unless I'm missing something.
Definitely worth having a closer look. Autograd via function reflection should be faster than backprop. And if it's running on AMD GPUs then it's quite intriguing.
14 comments
[ 0.24 ms ] story [ 46.0 ms ] threadEveryone knows JAX is what Google realised Tensorflow should have been when they realised how much of a joy Pytorch was to use. I actually think JAX does offer some advantages, not least true numpy interoperability. However, not mentioning *torch a single time in the blog post seems a little disingenuous for a Google-owned deep learning enterprise.
A technical reason for "why not pytorch" is that JAX was also built in part to expose and leverage the power of the XLA compiler, which is at least for the moment a pretty uniquely powerful tool for producing efficient, highly-scalable accelerator code.
I should underline that this is a friendly community of peers though: there is a lot of respect for Pytorch, which in turn was certainly influenced by the original Autograd that many of the JAX devs also worked on. JAX (and its fancier sibling Dex) beyond being useful tools are also still research projects in and of themselves seeking to advance our ideas on how to write expressive, powerful numerical code on modern architectures.
OTOH I'm not sure most people know what tasks are GPU-worthy or not. I haven't the slightest idea of why MCMC/Variational Bayes is amenable to GPU speedups and Persistent Homology isn't.
Yes, that is in active development: https://pymc-devs.medium.com/the-future-of-pymc3-or-theano-i...
JAX enables using (parts of) existing numpy codebases in disciplines other than deep learning. Autodiff and compilation to GPUs are very useful for all kinds of algorithms and processing pipelines.
I haven't found a way that I like. The closest I've come is tftorch: https://twitter.com/theshawwn/status/1311925180126511104
I think it's important to have global scope names. "biggan.discriminator.3.conv1.kernel.b" is perfectly sensible: it's the bias value of the kernel for the convolution of the third block of your biggan discriminator.
Everyone tries to treat model variables as interchangeable nameless parts. I hate it. Every variable has a global name, conceptually.
A global name also has other benefits. It becomes far easier to create EMA weights, for example, since you can filter the variables by name.
Their is no clear best pattern at the moment. My favorite would be Flax and their recent Linen API which is a refined effort that pays off when using their framework.
But, had a look in the code and jax has cublas and RoCm blas, and it looks like there is a flow where it uses the gpu directly, unless I'm missing something.
Definitely worth having a closer look. Autograd via function reflection should be faster than backprop. And if it's running on AMD GPUs then it's quite intriguing.