JAX LLM examples A collection (in progress) of example high-performance large language model implementations, written with JAX. Current contents include: DeepSeek R1 Llama 4 Llama 3 Qwen 3 For multi-host cluster setup and distributed training, see multi_host_README.md and the tpu_toolkit.sh script.