Chao Ma
April 25, 2026
Notes on the JAX ecosystem for machine learning, from the core transformation model to the surrounding training stack.
JAX AI Stack How JAX, XLA, Flax NNX, Optax, Orbax, and Grain fit together into a modern training stack, plus the role of jit, grad, and vmap.
---
title: "JAX"
author: "Chao Ma"
date: "2026-04-25"
---
Notes on the JAX ecosystem for machine learning, from the core transformation model to the surrounding training stack.
---
::: {.content-grid}
::: {.content-card}
**[JAX AI Stack](jax-ai-stack.qmd)**
How JAX, XLA, Flax NNX, Optax, Orbax, and Grain fit together into a modern training stack, plus the role of `jit`, `grad`, and `vmap`.
:::
:::