> For clean Markdown content of this page, append .md to this URL. For the complete documentation index, see https://docs.nvidia.com/jax-toolbox/llms.txt

# 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 | Models | Use cases | Container |
| :--- | :---: | :---: | :---: |
| [MaxText](/jax-toolbox/frameworks/max-text/overview) | GPT, LLaMA, Gemma, Mistral, Mixtral | pre-training | `ghcr.io/nvidia/jax:maxtext` |
| [AXLearn](/jax-toolbox/frameworks/ax-learn) | Fuji | pre-training | `ghcr.io/nvidia/jax:axlearn` |
| [AlphaFold3](https://github.com/google-deepmind/alphafold3) | EvoFormer | inference | `ghcr.io/nvidia/jax:alphafold` |

## Framework guides

- [MaxText](/jax-toolbox/frameworks/max-text/overview) — high-performance, scalable LLM
  framework by Google, including an [NVFP4 training example](/jax-toolbox/frameworks/max-text/nvfp-4-training-example).
- [AXLearn](/jax-toolbox/frameworks/ax-learn) — 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](/jax-toolbox/build-status).