The JAX AI Stack is an official meta-package that bundles the most commonly used JAX libraries. It's the easiest way to get started with a complete deep learning setup.
# Install the whole stack at once
# pip install jax-ai-stack
# This gives you:
# - jax (core)
# - flax (neural networks)
# - optax (optimizers)
# - orbax (checkpointing)
# - ml_dtypes (bfloat16, etc.)
# - grain (data loading)
# - chex (testing utilities)
import jax
from flax import nnx
import optax
import orbax.checkpoint as ocp
# Version check (early 2026):
# jax ~0.9.x, flax ~0.12.x, optax ~0.2.xPallas: Custom GPU/TPU Kernels
Pallas is JAX's kernel language for writing custom low-level operations that run directly on GPUs and TPUs. When XLA's automatic optimization isn't enough, Pallas lets you write hand-optimized kernels while staying in Python.
import jax
from jax.experimental import pallas as pl
import jax.numpy as jnp
# A simple Pallas kernel: vector addition
def add_kernel(x_ref, y_ref, o_ref):
"""Pallas kernel: reads from x and y, writes to o."""
o_ref[...] = x_ref[...] + y_ref[...]
# Launch the kernel
def pallas_add(x, y):
return pl.pallas_call(
add_kernel,
out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
)(x, y)
x = jnp.ones(1024)
y = jnp.ones(1024) * 2
result = pallas_add(x, y)
print(result[:5]) # [3. 3. 3. 3. 3.]Pallas translates to Triton on GPUs and Mosaic on TPUs. The name is a reference: in Greek mythology, Pallas is the parent of Triton.
💡 Why This Matters
Pallas is how Google implements things like FlashAttention and other custom fused kernels for JAX. Most users will never need Pallas directly — XLA's compiler handles 99% of optimization automatically. But if you're implementing novel attention mechanisms, custom quantization kernels, or other operations where XLA's fusion doesn't produce optimal code, Pallas gives you the escape hatch to write hand-optimized kernels while staying in the JAX ecosystem.