Pax is a programming framework for training very large neural network models across multiple TPU slices or pods. It builds on JAX and related libraries to support scalable model development and distributed training.
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.