For AI agents: a documentation index is available at the root level at /llms.txt. Append /llms.txt to any URL for a page-level index, or .md for the markdown version of any page.

Getting Started
Frameworks
- AXLearn
Performance & Profiling
Resiliency
- Resilient JAX with Ray
- Example Walkthrough
Inference
- JAX-vLLM Rollout Offloading Bridge
Customization
- GPU Kernels
Reference

Introduction
JAX on NVIDIA GPU Stack
Frameworks & Supported Models
Build Pipeline Status
Environment Variables
Overview
NVFP4 Training Example
AXLearn
Profiling
GPU Performance
PGLE (Profile-Guided Latency Estimation)
Native FP8
nsys-jax
Resilient JAX with Ray
Example Walkthrough
JAX-vLLM Rollout Offloading Bridge
GPU Kernels
Container Versions
Staging Containers
Cloud Providers
FAQ
Resources & Videos

On this page

Resources
Videos

Resources & Videos

External references, talks, and related material

||View as Markdown|

Resources

JAX | NVIDIA NGC Container
Slurm and OpenMPI zero config integration
Adding custom GPU ops
Triaging regressions

Videos

Equinox for JAX: The Foundation of an Ecosystem for Science and Machine Learning
Scaling Grok with JAX and H100
JAX Supercharged on GPUs: High Performance LLMs with JAX and OpenXLA
What’s New in JAX | GTC Spring 2024
What’s New in JAX | GTC Spring 2023

FAQ

Developer-friendly docs for your API

Privacy Policy | Your Privacy Choices | Terms of Service | Accessibility | Corporate Policies | Product Security | Contact

Copyright © 2026, NVIDIA Corporation.