Frameworks & Supported Models

We support and test the following JAX frameworks and model architectures. More details about each model and available containers can be found in their respective guides.

Framework guides

• MaxText — high-performance, scalable LLM framework by Google, including an NVFP4 training example.
• AXLearn — Apple’s deep-learning design framework built on JAX and XLA.

For the full list of published containers and their build/test status, see Build Pipeline Status.