Machine Learning with JAX - From Hero to HeroPro+ | Tutorial #2
❤️ Become The AI Epiphany Patreon ❤️ https://www.patreon.com/theaiepiphany 👨👩👧👦 Join our Discord community 👨👩👧👦 https://discord.gg/peBrCpheKE This is the second video in the JAX series of tutorials. JAX is a powerful and increasingly more popular ML library built by the Google Research team. The 2 most popular deep learning frameworks built on top of JAX are Haiku (DeepMInd) and Flax (Google Research). In this video, we continue on and learn additional components needed to train complex ML models (such as NNs) on multiple machines! ▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ ✅ My GitHub repo: https://github.com/gordicaleksa/get-started-with-JAX ✅ JAX GitHub: https://github.com/google/jax ✅ JAX docs: https://jax.readthedocs.io/ ▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ ⌚️ Timetable: 00:00:00 My get started with JAX repo 00:01:25 Stateful to stateless conversion 00:11:00 PyTrees in depth 00:17:45 Training an MLP in pure JAX 00:27:30 Custom PyTrees 00:32:55 Parallelism in JAX (TPUs example) 00:40:05 Communication between devices 00:46:05 value_and_grad and has_aux 00:48:45 Training an ML model on multiple machines 00:58:50 stop grad, per example grads 01:06:45 Implementing MAML in 3 lines 01:08:35 Outro ▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 💰 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ć Bartłomiej Danek Zvonimir Sabljic ▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 💼 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/ ▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ #jax #machinelearning #framework
