Overview

JAX is a Python library for high-performance machine learning research. It combines the familiarity of NumPy with hardware acceleration and composable function transformations. Google researchers have built and trained models like Gemini and Gemma on JAX, and it’s also used by researchers for a wide range of advanced applications. This talk will provide an introduction to JAX and the Flax neural network library, showcasing recent features and how to get started.