Repolode
Submit
jax-ml

jax-ml/jax-ai-stack

A pip-installable bundle that pins a tested set of JAX libraries (Flax, Optax, Orbax, Chex, Grain) for AI and numerical computing.

View on GitHub
Stars
315
Forks
51
Open issues
38
Last commit
25 days ago

What is it?

JAX AI Stack is a single entry point for installing a set of JAX-based libraries that work together. Running pip install jax-ai-stack pulls in JAX along with component projects such as Flax for building neural networks, Optax for gradient processing and optimization, Orbax for checkpointing, Chex for writing reliable JAX code, Grain for data loading, and ml_dtypes for NumPy dtype extensions. The versions are pinned to combinations verified by the repository's integration tests.

JAX itself is an array computation library with program transformations such as jit, vmap and grad. It contains little that is specific to neural networks, and the surrounding ecosystem is developed as separate packages. This stack collects the ones Google developers use for generative AI models, so users don't have to resolve compatible versions by hand.

Optional extras add TensorFlow Datasets, and hardware-specific JAX builds for CUDA GPUs or TPUs can be installed in the same command.

Key features

  • 01Installs JAX with Flax, Optax, Orbax, Chex, Grain and ml_dtypes at versions tested together.
  • 02Provides JAX transformations such as jit, vmap and grad.
  • 03Offers an optional tensorflow-datasets extra via pip.
  • 04Supports hardware-specific JAX installs for CUDA GPUs and TPUs.

Use cases

  • Setting up a JAX environment for training neural networks with Flax and Optax.
  • Saving and restoring model checkpoints with Orbax in JAX projects.
  • Running JAX workloads on NVIDIA GPUs or Google TPUs with a known-good dependency set.

Quick start

  1. 1

    Install the JAX AI stack from PyPI

    pip install jax-ai-stack
  2. 2

    Optionally add TensorFlow Datasets support via the tfds extra

    pip install jax-ai-stack[tfds]
  3. 3

    Optionally install with GPU support, adding the JAX CUDA extra in the same command

    pip install jax-ai-stack "jax[cuda]"

Things to know

  • This is a bundle of other libraries, so it offers no standalone functionality beyond the packages it pins.
  • The README says documentation and tutorials are still a work in progress.

In collections

Run and train AI on GPUs

Topics

#jax#machine-learning#deep-learning#python#numerical-computing#gpu#tpu

FAQ

What is jax-ai-stack?
A pip-installable bundle that pins a tested set of JAX libraries (Flax, Optax, Orbax, Chex, Grain) for AI and numerical computing. JAX AI Stack is a single entry point for installing a set of JAX-based libraries that work together. Running pip install jax-ai-stack pulls in JAX along with component projects such as Flax for building neural networks, Optax for gradient processing and optimization, Orbax for checkpointing, Chex for writing reliable JAX code, Grain for data loading, and ml_dtypes for NumPy dtype extensions. The versions are pinned to combinations verified by the repository's integration tests.
Is jax-ai-stack free to use? What license does it have?
jax-ai-stack is open source under the Apache-2.0 license. Check the license terms before using it in commercial products.
Is jax-ai-stack still maintained?
Actively maintained. The last commit to jax-ai-stack was on Oct 13, 2026.
What can I use jax-ai-stack for?
Setting up a JAX environment for training neural networks with Flax and Optax. Saving and restoring model checkpoints with Orbax in JAX projects. Running JAX workloads on NVIDIA GPUs or Google TPUs with a known-good dependency set.
How do I install jax-ai-stack?
Install the JAX AI stack from PyPI: pip install jax-ai-stack Optionally add TensorFlow Datasets support via the tfds extra: pip install jax-ai-stack[tfds] Optionally install with GPU support, adding the JAX CUDA extra in the same command: pip install jax-ai-stack "jax[cuda]"
Where can I get jax-ai-stack?
The source code and installation instructions are on GitHub: https://github.com/jax-ml/jax-ai-stack

Summary generated by AI from the project's README. GitHub stats refresh daily.

Related projects