JAX AI Stack

JAX is a Python package for array-oriented computation and program transformation. Built around it is a growing ecosystem of packages for specialized numerical computing across a range of domains; an up-to-date list of such projects can be found at Awesome JAX.

Though JAX is often compared to neural network libraries like PyTorch, the JAX core package itself contains very little that is specific to neural network models. Instead, JAX encourages modularity, where domain-specific libraries are developed separately from the core package: this helps drive innovation as researchers and other users explore what is possible.

Within this larger, distributed ecosystem, there are a number of projects that Google researchers and engineers have found useful for implementing and deploying the models behind generative AI tools like Imagen, Gemini, and more. The JAX AI stack serves as a single point-of-entry for this suite of libraries, so you can install and begin using many of the same open source packages that Google developers are using in their everyday work.

To get started with the JAX AI stack, you can check out Getting started with JAX. This is still a work-in-progress, please check back for more documentation and tutorials in the coming weeks!

Installing the stack

The stack can be installed with the following command:

pip install jax-ai-stack

This pins particular versions of component projects which are known to work correctly together via the integration tests in this repository. Packages include:

• JAX: the core JAX package, which includes array operations and program transformations like jit, vmap, grad, etc.
• flax: build neural networks with JAX
• ml_dtypes: NumPy dtype extensions for machine learning.
• optax: gradient processing and optimization in JAX.
• orbax: checkpointing and persistence utilities for JAX.
• chex: utilities for writing reliable JAX code.
• grain: data loading.

Optional packages

Additionally, there are optional packages you can install with pip extras.

The following command:

pip install jax-ai-stack[tfds]

will install a compatible version of tensorflow and tensorflow-datasets.

Hardware support

To install jax-ai-stack with hardware-specific JAX support, add the JAX installation command in the same pip install invocation. For example:

pip install jax-ai-stack "jax[cuda]"  # JAX + AI stack with GPU/CUDA support
pip install jax-ai-stack "jax[tpu]"  # JAX + AI stack with TPU support

For more information on available options for hardware-specific JAX installation, refer to JAX installation.