This talk is for anyone interested in building and understanding modern machine learning models. We'll start with a quick introduction to JAX, highlighting its core principles like automatic differentiation and JIT compilation, which make it so powerful for high-performance computing.

We'll then dive into Flax, a neural network library built on top of JAX. You'll learn how Flax simplifies model building with its flexible module system. Instead of just talking about it, we'll get hands-on.

The second half of the session is a live coding tutorial where we'll recreate a research paper's model from scratch. Attendees will follow along, gaining practical experience in using JAX and Flax to implement a real-world architecture. By the end, you'll not only understand the "what" and "why" behind JAX and Flax but also have the skills to start building and replicating your own research models.