JAX is a Python library built for high-performance machine-learning research that leverages the familiarity of NumPy with the power of hardware acceleration on CPUs, GPUs, and TPUs. Its composable function transformations allow you to seamlessly combine automatic differentiation, automatic batching, end-to-end compilation, and both data and model parallelism.

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, from studying training dynamics of neural networks, to probabilistic programming, to scientific applications in physics and biology.

This jam-packed talk will show you what it's like to use JAX and the Flax neural network library, using live coding examples. You’ll learn about how it’s powering new research, and how you can get started using this fast, flexible, and fun tool.

*This kick-off event will be happening at Google for Korea 2025.