Machine Learning with JAX - From Zero to Hero | Tutorial #1
❤️ Become The AI Epiphany Patreon ❤️ https://www.patreon.com/theaiepiphany 👨👩👧👦 Join our Discord community 👨👩👧👦 https://discord.gg/peBrCpheKE With this video I'm kicking off a series of tutorials on JAX! 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 I cover the basics as well as the nitty-gritty details of jit, grad, vmap, and various other idiosyncrasies of JAX. ▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ ✅ JAX GitHub: https://github.com/google/jax ✅ JAX docs: https://jax.readthedocs.io/ ✅ My notebook: https://github.com/gordicaleksa/get-started-with-JAX ✅ Useful video on autodiff: https://www.youtube.com/watch?v=wG_nF1awSSY ▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ ⌚️ Timetable: 00:00:00 What is JAX? JAX ecosystem 00:03:35 JAX basics 00:10:05 JAX is accelerator agnostic 00:15:00 jit explained 00:17:45 grad explained 00:27:25 The power of JAX autodiff (Hessians and beyond) 00:31:00 vmap explained 00:36:50 JAX API (NumPy, lax, XLA) 00:39:40 The nitty-gritty details of jit 00:46:55 Static arguments 00:50:05 Gotcha 1: Pure functions 00:56:00 Gotcha 2: In-Place Updates 00:57:35 Gotcha 3: Out-of-Bounds Indexing 00:59:55 Gotcha 4: Non-Array Inputs 01:01:50 Gotcha 5: Random Numbers 01:09:40 Gotcha 6: Control Flow 01:13:45 Gotcha 7: NaNs and float32 02:15:25 Quick summary 02:16:00 Conclusion: who should be using JAX? 02:17:10 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 ▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬ 📄 Website - https://gordicaleksa.com/ 💼 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
