JAX is a Python library for accelerator-oriented array computation and program transformation, designed for numerical computing and large-scale machine learning. It transforms Python and NumPy-style functions through automatic differentiation with `jax.grad`, automatic batching and vectorization with `jax.vmap`, and just-in-time compilation with `jax.jit`. It uses XLA to compile computations for GPUs, TPUs, and other accelerators. JAX supports compiler-driven multi-device execution: `Mesh`, `PartitionSpec`, and `NamedSharding` declare array placement across device meshes, allowing distributed training code to include automatically managed parallel work such as gradient synchronization. Its compilation tools also include ahead-of-time compilation, portable artifact export with `jax.export`, and conversion to TensorFlow SavedModels through `jax2tf` for serving integrations. The project describes itself as a research project rather than an official Google product.
XLA (Accelerated Linear Algebra) is an open-source machine-learning compiler for GPUs, CPUs, and other ML accelerators. It takes models from frameworks such as JAX, PyTorch, and TensorFlow, traces and compiles them into optimized executable code for the target hardware. The project is developed under OpenXLA and is primarily used through the documentation and integration of the corresponding ML framework; its repository is intended mainly for compiler contributors and frontend or hardware-backend integrators.
Searchable transcript of Getting started with JAX on NVIDIA GPUs — Google Cloud Tech (04:00). Search for a phrase, then click its timestamp to jump straight to that moment in the video.
Captions sourced from the original video on YouTube, published by Google Cloud Tech. The video, its captions and all related intellectual property remain the property of their respective owners; AINotes claims no ownership. Provided for research, accessibility and search — see the Transcript Notice and Copyright Policy.
00:00 Hi there. Welcome to our course on JAX AI on GPU. I'm Ivan. I'm an AI developer engineer at Google Cloud. >> And I'm Katya, a developer advocate at Nvidia. >> In this first session, we will cover how to get started with JAX on GPU. And we will provide some insights about JAX compilation as well as how to profile and debug. Ready, Katya? Let's go. >> Before we optimize anything, let's investigate.
00:26 The GPU looks powerful, the notebook looks up, and somehow the code is still slow. First question, is JAX even using the GPU? >> The Python code you write is only the top layer. JAX AI traces it, XLA compiles it, and the Nvidia stack runs it. If we know which layer we are looking at, we won't be surprised by any expected behavior. >> Start at the bottom, run nvidia-smi, then ask JAX directly with JAX devices and JAX default backend.
00:57 If you see only CPU devices, stop. We need to see that the GPU is available. >> Once the GPU is a visible, JAX feels like NumPy. But now location matters. A NumPy array lives on the host, and JAX array can live on the GPU, the so-called device. So, the useful debugging question becomes, where does this value live? >> Tiny warning label, and float loss are fine for checks and logs.
01:28 Inside a hot loop, they can pull data back to Python and force the device to wait. That is why if you're bringing your code to production, it is important to make sure that there are no associated leaks. >> That's why you need to know the three beginner primitives in JAX, which are jit, grad, and vmap, which are functional names for compiling, differentiating, and batching.
01:50 You write an ordinary array functions, then transform them. And the one that explained the monster prices is JIT. At this point, you know that you can compile with JIT. But JAX works differently compared to other frameworks. This also means that you need to know about how compilation works. Let's pause and guess. First call has shape four. Second call has shape four again.
02:14 Third call has shape five. [music] Which ones compile? Same shape and dtype usually reuse the executable. A new signature may compile again. This is why fixed shape is matter. Right batches, variable sequence lens, and more final batches can accidentally create a new compiled programs. Padding plus a mask often keeps the shape stable. Code can also look normal, but fail under tracing.
02:44 Inside JIT, array values are traced values. Python cannot branch on them like ordinary numbers. Use JNP.where, lax.cond, and lax.scan for JAX-friendly control flow. So, there are many things to consider when diagnosing unexpected behavior [music] on JAX. Exactly, Katya. And timing has its own trap. JAX launches GPU work asynchronously. If your timer does not block, you may measure how fast Python asked the GPU to work, not how long the GPU worked.
03:16 Think for a moment. Which timing number would you trust? If timing is still strange, profile once before guessing. Look for compilation, transfers, empty GPU gaps, many tiny kernels, batch size effects, and memory pressure. These are the things affecting performance you want to keep a closer attention at. So, the question is not why is JAX slow. It is are we on GPU, compiling, transferring, changing shape, or underfilling the device? Next, we apply the checklist to a real training loop.