Implementing a deep learning model often feels like working in a chaotic kitchen—a frantic scramble of complex code, unpredictable performance, and seeds that must be 42.
But what if we have a "mise en place" system for implementing AI, where every function is a perfectly prepared ingredient, ready to be combined with precision and speed. It’s the disciplined approach that JAX leverages to unlock its incredible performance.
In this session, we'll put on our aprons to cook an image diffusion model from scratch. You'll learn how the @jax.jit decorator acts as the "Fire!" command that executes our entire recipe at incredible speed, and how @jax.vmap lets us shout "Corner!" to serve a whole batch of images in parallel, without a single messy collision.
This talk is for any Python developer who wants to roll up their sleeves and truly understand how Denoise Diffusion Models (a.k.a. a Score-Based Generative Model), the backbone of the popular text-to-image and text-to-video work. This will be a step by step recipe dissecting the ingredients that make these models.
No prior culinary or JAX expertise is required, but you will leave the session wanting to say "Yes, Chef!" to home cook your next model, or at least to understand how it has been cooked.