Using JAX to Accelerate Research

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".

Can you explain how JAX compares to pytorch? AFAIK pytorch also closely resembles numpy API.

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.

> why not Pytorch?

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 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.

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.

Instead of.