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

Choose to Build with AI
Matched to Google

AI Maker Residence 3

The third AI workshop taught by our legendary teacher, Nick Sarafa. This is a full-day hands-on training workshop for purposeful co-creation with AI using Claude Code. Imagine having access to hundreds of billions of dollars of computing power and knowing exactly how to make it work for you through the power of super intelligence.

◆ Fri 09 Oct 2026 ◆ KOKO Cafe, London ◆ With Nick Sarafa
AI Maker Residence 3
Live event
AI Maker Residence 3
Fri 09 Oct 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.