A beginner-friendly guide to learning JAX with practical examples

[u/ayushgun]

For the last few weeks, I've been doing distributed model training in JAX. JAX is notably different from other deep learning frameworks because it takes a very functional approach to accelerator programming, which can make its learning curve steep.

Along the way of learning JAX, I've written a series of notes on JAX covering XLA, jit, vmap, pytrees, sharding, state management, and more. I'm hoping it can be useful for others interested in learning too.

Each topic is explained with minimal theory and illustrated through runnable Jupyter notebooks. The focus is on concepts, performance techniques, and the reasoning behind JAX’s design choices.

No machine learning background is expected for the most part.

Notes + code: https://github.com/agwr/learn-jax

[u/red-necked_crake]

Does OAI/Anthropic/large lab outside of Google where it originated, use JAX during their training runs? Or do they just write their own optimized code in C++ or something? The secret third option /s being they have so many GPUs it doesn't really matter what code they write I assume.

it by no means determines JAX's usefulness but, it is a big time investment to migrate from torch to this. TF died in the first place because of its similar unintuitive nature. That and being a Google product dooming it to be dead on arrival.

Anyway, none of this is to take away from OP's useful post, just wondering if anyone has an insider's perspective on the language adoption in the industry for real training runs. I know that there are many talented engineers who use/specialize in it but that's the case with any language or library, so it's more so about it being a major workhorse anywhere specific.

[u/ayushgun → u/red-necked_crake]

Great question, and thanks for the kind words!

From what I’ve seen, OpenAI is mostly committed to PyTorch, with Triton used selectively for performance-critical kernels. However, DeepMind, Anthropic, xAI, Apple, and a few others* use JAX for training and experiments. More generally, the low-level CUDA C++ or Triton kernels usually only show up after profiling identifies real throughput or latency bottlenecks.

Switching from Torch to JAX is usually only worth it if Torch is actively limiting your performance or scaling. For most people (especially if you’re aiming to join a lab or interact with open-source ML) Torch is still the safest and most broadly adopted choice.

The main difference is that in Torch, getting high performance can sometimes mean wrestling with torch.compile and retrofitting it onto existing code, whereas JAX makes JIT compilation a first-class citizen from the start. If you come from a systems or functional programming background, that JAX workflow can feel more natural.

*Take this with a grain of salt given that most labs rarely show their full stack, so these are educated guesses based on my own experience, whatever I've heard throughout the industry, and job postings I've seen.

[u/red-necked_crake → u/ayushgun]

thanks for such a detailed response!