A programming framework designed for training large-scale neural network models so large that they span multiple TPU accelerator chip slices or pods. Pax is built on Flax, which is built on JAX. !Diagram indicating Pax's position in the software stack. Pax is built on top of JAX. Pax itself consists of three layers. The bottom layer contains TensorStore and Flax. The middle layer contains Optax and Flaxformer. The top layer contains Praxis Modeling Library. Fiddle is built on top of Pax.