JAX Container Early Access

JAX is a library for high-performance numerical computing and machine learning research.

JAX NGC containers come with all dependencies included, providing an easy place to start developing applications in areas such as NLP, computer vision, multimodality, physics-based simulations, reinforcement learning, drug discovery, and neural rendering.

For building neural networks, the JAX NGC Container includes Flax, a neural network library. We also include a container for Paxml, a framework for training LLMs such as GPT, and a container for T5x, a framework for training T5 and other Flax-based models. You can use the JAX, Paxml, or T5x containers for your deep learning workloads or install your favorite libraries on top.

JAX containers are now generally available on NGC.