Machine Learning with JAX - From Zero to Hero | Tutorial #1

❤️ Become The AI Epiphany Patreon ❤️

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

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.

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

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

⌚️ 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!

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

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

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

#jax #machinelearning #framework

[billykotsos]

This channel is going on the GOAT level status

[matthewkhoo]

When I started learning JAX, I personally think it stands for JIT (J), Autograd (A), XLA (X) which is essentially an abbreviation for a bunch of abbreviations. Given that those features are the 'highlights' of JAX, its very possible. If that's the case, pretty cool naming from DeepMind. Anyways, there aren't many comprehensive resources for JAX right now, so I'm really looking forward to this series! Cheers Aleksa.

[not_a_human_being]

50:40 - I would argue here, that it's not necessary to pass all the parameters into the function, as long as it's not changing any of the params, it's ok to use external globals(), like for some reference tables etc. This definition (though academically thorough), make practical application a bit more cumbersome. I believe that the better way to think "2." is sufficient to make this work. No need to pass long list of params. Just make sure not to update/change anything external inside the function, and whatever is not passed in is static. Alternatively, you can have "get jit_function" every time you anticipate that your globals might've changed. So, you will be effectively re-creating your jit function with new globals(). In some cases that feels much preferable to passing everything in. For instance, you can use all sorts of globals inside it, then just re-create it just before your training loop.

[akashraut]

Finally some jax tutorial.. Keep them coming

[abbashaider]

That is a great hands-on tutorial which has a perfect mix of theory and practical usage. Thanks

[wryltxw]

this video is such a great service to the community. really great examples to help better understand Jax at a nuanced level.

[mariuskombou]

Thanks for this video!! that was really interesting for a new user of JAX like me

[vijayanand]

Great video!🔥 Would need Paper implementations too💯

[sacramentofwilderness]

Thanks for the great tutorial with pointing out the strong and weak points of the JAX framework, with caveats and salient features. What makes me somehow confused -the behavior, that overshooting the index array clips the index to maximum or does nothing. In C/C++ if one does this usually if the displacement is small - some memory part outside the given data is modified, and for strongly index mistake one would receive SEGMENTATION FAULT. Clipping the index makes the program safer, but in addition to counterintuitive behavior is adds some small additional cost for fetching the index.

[carlosrondonmoreno]

Wonderful tutorial! Really really good examples. Thank you so much!

[johanngerberding]

congrats to your deepmind job man (read your post), nice channel, keep going!

[mikesmith]

Awesome work!! JAX is a fantastic library. This series is the reason I finally subscribed to your channel. Thanks for your work!

[bionhoward]

i love jax ... thank you for your work!

[TheParkitny]

great video and content, this channel needs more recognition.

[sarahel-sherif]

Great material and great efforts, excited to see FLAX and Haiku, thank you

[yulanliu]

JAX = Just After eXecution (related to the tracing behaviour)
JIT = Just In Time (related to the compilation behaviour)

[jawadmansoor]

Great video thanks, kindly complete the tutorial series on Flax as well.

[alexanderchernyavskiy]

oh, quite impressive series with the perfect explanation

[PhucLe-qsnx]

JAX is Just After eXecution, represent the paradigm of tracing and transform (grad, vmap, jit, ..) after the first execution.

[sallanmega]

Thank you for the amazing content.

Greetings from Spain