Uploaded October 2021 | Updated September 2026, 1 week ago
β€οΈ Become The AI Epiphany Patreon β€οΈ
patreon.com/theaiepiphany
π¨βπ©βπ§βπ¦ Join our Discord community π¨βπ©βπ§βπ¦
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: github.com/google/jax
β JAX docs: jax.readthedocs.io
β My notebook: github.com/gordicaleksa/get-started-with-JAX
β Useful video on autodiff: 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 - patreon.com/theaiepiphany
One-time donation - paypal.com/paypalme/theaiepiphany
Huge thank you to these AI Epiphany patreons:
Eli Mahler
Petar VeliΔkoviΔ
BartΕomiej Danek
Zvonimir Sabljic
β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬
π Website - gordicaleksa.com
πΌ LinkedIn - linkedin.com/in/aleksagordic
π¦ Twitter - twitter.com/gordic_aleksa
π¨βπ©βπ§βπ¦ Discord - discord.gg/peBrCpheKE
πΊ YouTube - youtube.com/c/TheAIEpiphany
π Medium - gordicaleksa.medium.com
π» GitHub - github.com/gordicaleksa
π’ AI Newsletter - aiepiphany.substack.com
β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬
#jax #machinelearning #framework
β€οΈ Become The AI Epiphany Patreon β€οΈ
patreon.com/theaiepiphany
π¨βπ©βπ§βπ¦ Join our Discord community π¨βπ©βπ§βπ¦
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: github.com/google/jax
β JAX docs: jax.readthedocs.io
β My notebook: github.com/gordicaleksa/get-started-with-JAX
β Useful video on autodiff: 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 - patreon.com/theaiepiphany
One-time donation - paypal.com/paypalme/theaiepiphany
Huge thank you to these AI Epiphany patreons:
Eli Mahler
Petar VeliΔkoviΔ
BartΕomiej Danek
Zvonimir Sabljic
β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬
π Website - gordicaleksa.com
πΌ LinkedIn - linkedin.com/in/aleksagordic
π¦ Twitter - twitter.com/gordic_aleksa
π¨βπ©βπ§βπ¦ Discord - discord.gg/peBrCpheKE
πΊ YouTube - youtube.com/c/TheAIEpiphany
π Medium - gordicaleksa.medium.com
π» GitHub - github.com/gordicaleksa
π’ AI Newsletter - aiepiphany.substack.com
β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬β¬
#jax #machinelearning #framework
