Jax

Stack

• Flax: model authoring

• Optax: composable optimizers

• Orbax large-scale asynchronous checkpointing

From

• Deepmind

• Midjourney

• Cohere

• Infrastructure:
  XLA (operation fusion & memory optimization), Pathways (tens of thousands of chips distributed & fault recovery).

• Advanced Tools: Pallas/
  tokamax

  openxla • Updated 2026 Jan 6 0:34
  (custom & optimal kernels),
  Qwix (non-invasive quantization),
  Jax Grain (reproducible data pipelines).

• End-to-End: MaxText/MaxDiffusion (pre-training), Tunix (FT), vLLM JAX TPU backend.

JAX AI Stack — JAX AI Stack

JAX AI Stack tested releases provide high reliability by ensuring compatibility across its libraries

https://jaxstack.ai/

Google for Developers Blog - News about Web, Mobile, AI and Cloud

Introducing the JAX AI Stack: a modular, end-to-end, production-ready platform built on JAX and co-designed with Cloud TPUs for scalable ML.

https://developers.googleblog.com/building-production-ai-on-google-cloud-tpus-with-jax/

Kakao’s journey with JAX and Cloud TPUs | Google Cloud Blog

Kakao’s approach provides a compelling example of the high-performance array computing framework JAX for AI model development at scale.

https://cloud.google.com/blog/products/infrastructure-modernization/kakaos-journey-with-jax-and-cloud-tpus

Programming

Programming TPUs in JAX | How To Scale Your Model

How to use JAX to program TPUs efficiently! Much of this section is taken from <a href='https://jax.readthedocs.io/en/latest/jep/14273-shard-map.html'>here</a>.

https://jax-ml.github.io/scaling-book/jax-stuff/

Jax for Pytorch

A guide to JAX for PyTorch developers | Google Cloud Blog

PyTorch users can learn about JAX in this tutorial that connects JAX concepts to the PyTorch building blocks that they’re already familiar with.

https://cloud.google.com/blog/products/ai-machine-learning/guide-to-jax-for-pytorch-developers