Introduction to JAX: High-Performance Machine Learning Library

This 17-minute Google I/O 2025 talk introduces JAX, a Python library for high-performance machine learning research that combines NumPy's familiarity with hardware acceleration and composable function transformations. Learn how Google researchers have built and trained models like Gemini and Gemma using JAX, and discover how researchers apply it across various advanced applications. Speakers Yufeng Guo and Robert Dyro provide an introduction to JAX and the Flax neural network library, highlighting recent features and offering guidance on getting started with these powerful tools for AI development.