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.
StableHLO is a backward-compatible operation set for high-level machine-learning operations and a portability layer between ML frameworks and compilers. Frameworks such as TensorFlow, JAX, and PyTorch can produce StableHLO programs, while compilers such as XLA and IREE can consume them; the programs can also be used for on-device execution through Google AI Edge. Based on the MHLO dialect, StableHLO adds serialization and versioning through MLIR bytecode, with backward- and forward-compatibility guarantees. The project includes the StableHLO specification and an MLIR-based C++ and Python implementation for defining programs and integrating them with compatible compilers.
TensorFlow is an open-source, end-to-end machine-learning platform originally developed by Google Brain. It provides tools, libraries, community resources, and stable Python and C++ APIs for building, training, and deploying machine-learning applications, with support for CPU and GPU execution and additional device plugins. In TensorFlow-based deployment stacks, it provides serving infrastructure and the SavedModel format; JAX graphs can be converted to SavedModel for integration with TensorFlow serving. The project is distributed through PyPI packages, Docker containers, and source builds, including CPU-only and nightly packages.
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 From notebook to production: Serving JAX at scale — Google Cloud Tech (03:05). 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 A model that works in a notebook can still underperform when you serve it. The first user asks for a prediction and the server goes into a long compilation who pays the compile cost. >> In the last video, we saved a train transformer checkpoint. Now we rebuild the model structure, load the orbox parameters and place them on the serving device. >> There are different patterns you can consider.
00:25 Jax G for a Python service. Hi head of time compile for startup time compilation. Jax export for portable Jax artifacts and JAX 2TF when the serving stack is TensorFlow based. >> First Jaxjit is the simplest in process pass. It's great for Python service but the first call may compile. Consider warming it up before real traffic. Second option is ahead of time compilation which makes the stage explicit.
00:58 Lower produce stable HLO and compile produces the executable. Then server calls execute without the first call compile spike. >> Stable HLO is not a new model. It's the lowest program the compiler sees. You inspect it when you need to debug or confirm what got staged. >> [music] >> Batch size is still a thing to consider. New shapes may compile separately, but larger batches can improve tokens per second by a mortisizing overhead.
01:29 Same shape lesson we share before, but now at a second time. Third option is Jax export which serializes a stable HLO based artifact. The notebook writes it, reads it back, calls it and checks that predictions match. If your platform speaks Jax runtime, this is the portable route. >> Fourth and last option is Jax 2TF for TensorFlow infrastructure. If the serving stack expects a save model, convert with native serialization and save the standard TF artifact.
02:09 To wrap up, choose by deployment target, Python service, AOTJ runtime, portable JAX artifact or TensorFlow serving stack. The checkpoint is the same but the serving wrapper changes. >> So to conclude, in the whole course we had one mental model J traces, XLA compiles and the Nvidia stack runs. When something surprises you, keep in mind these four layers and start your investigation.
02:40 >> Don't forget to verify the GPU stabilized [music] shapes. Time is blocking profile before guessing and avoid accidental host transfers. >> And that's it. This is the path from your first GPU chuck to serving a Jax model on GPU. >> Now go make your GPU busy.