Machine Learning with Flax - From Zero to Hero

From the creator

❤️ Become The AI Epiphany Patreon ❤️ https://www.patreon.com/theaiepiphany 👨‍👩‍👧‍👦 Join our Discord community 👨‍👩‍👧‍👦 https://discord.gg/peBrCpheKE ❤️ Deepnote referral link: https://deepnote.com/referral?token=823d18856ad5 - get 20h of free "Pro machine" + support the channel at 0 cost to you! ❤️ In this video I cover Flax - a JAX-based machine learning library. It's a part of my machine learning with JAX series of videos! You'll learn everything you need to get you started building ML models in Flax. ▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ ✅ Flax notebook (do follow along!): https://github.com/gordicaleksa/get-started-with-JAX/blob/main/Tutorial_4_Flax_Zero2Hero_Colab.ipynb ✅ Flax GitHub: https://github.com/google/flax ✅ Flax docs: https://flax.readthedocs.io/en/latest/ ✅ HuggingFace Flax community week: https://github.com/huggingface/transformers/tree/master/examples/research_projects/jax-projects ▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ ⌚️ Timetable: 00:00:00 Intro - Flax is performant and reproducible 00:02:00 Deepnote walk-through (sponsored) 00:04:57 Flax basics 00:14:05 Flax vs Haiku 00:17:25 Benchmarking Flax 00:18:28 Linear regression toy example 00:27:05 Introducing Optax (Adam state example) 00:33:15 Creating custom models 00:39:30 self.param example 00:45:55 self.variable example 00:57:15 Handling dropout, BatchNorm, etc. 01:02:25 CNN on MNIST example 01:11:00 TrainState source code 01:13:00 CNN dropout modification 01:15:05 Outro and summary ▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 💰 BECOME A PATREON OF THE AI EPIPHANY ❤️ If these videos, GitHub projects, and blogs help you, consider helping me out by supporting me on Patreon! The AI Epiphany - https://www.patreon.com/theaiepiphany One-time donation - https://www.paypal.com/paypalme/theaiepiphany Huge thank you to these AI Epiphany patreons: Eli Mahler Petar Veličković ▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 💼 LinkedIn - https://www.linkedin.com/in/aleksagordic/ 🐦 Twitter - https://twitter.com/gordic_aleksa 👨‍👩‍👧‍👦 Discord - https://discord.gg/peBrCpheKE 📺 YouTube - https://www.youtube.com/c/TheAIEpiphany/ 📚 Medium - https://gordicaleksa.medium.com/ 💻 GitHub - https://github.com/gordicaleksa 📢 AI Newsletter - https://aiepiphany.substack.com/ ▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ #flax #jax #deeplearning

Stop watching · start building
Matched to Google

Makers residency 4th edition

Your idea can be anything. At our last workshops, people built plant-based nutrition businesses, a way to measure planetary resilience, crochet headwear and a way to connect lonely people living in the same building. An asset manager went on to raise an additional £25M on their fund within nine months. The fourth workshop is a full day at KOKO with Nick Sarafa, learning to build with Claude Code. Bring the thing you keep meaning to start. We’ll bring the pizza.

◆ Fri 13 Nov 2026 ◆ KOKO, London
Makers residency 4th edition
Live event
Makers residency 4th edition
Fri 13 Nov 2026

More like this

Running one yourself?

List your AI event,
wherever it is.

A meetup, a workshop, a hackathon, a conference. Any city, or online. Tell us about it and it lands in front of people already learning this stuff.