❤️ Become The AI Epiphany Patreon ❤️
  / theaiepiphany  

👨‍👩‍👧‍👦 Join our Discord community 👨‍👩‍👧‍👦
  / discord  

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-s...

✅ 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 -   / theaiepiphany  
One-time donation - https://www.paypal.com/paypalme/theai...

Huge thank you to these AI Epiphany patreons:
Eli Mahler
Petar Veličković
Bartłomiej Danek
Zvonimir Sabljic

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬

💼 LinkedIn -   / aleksagordic  
🐦 Twitter -   / gordic_aleksa  
👨‍👩‍👧‍👦 Discord -   / discord  

📺 YouTube -    / theaiepiphany  
📚 Medium -   / gordicaleksa  
💻 GitHub - https://github.com/gordicaleksa
📢 AI Newsletter - https://aiepiphany.substack.com/

▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬

#jax #machinelearning #framework