Home

Awesome

Get started with JAX! :computer: :zap:

The goal of this repo is to make it easier to get started with JAX, Flax, and Haiku!

JAX ecosystem is becoming an increasingly popular alternative to PyTorch and TensorFlow. :sunglasses:

<br/> <br/> <p align="center"> <img src="readme_pics/jax_logo.png" width="300"/> </p> <br/> <br/>

Note: I'm only going to recommend content that I've personally analyzed and found useful here. If you want a comprehensive list check out the awesome-jax repo.

Table of Contents

My Machine Learning with JAX Tutorials

Tip on how to use notebooks: just open the notebook directly in Google Colab (you'll see a button on top of the Jupyter file which will direct you to Colab). This way you can avoid having to setup the Python env! (This was especially convenient for me since I'm on Windows which is still not supported)

Tutorial #1: From Zero to Hero

In this video, we start from the basics and then gradually dig into the nitty-gritty details of jit, grad, vmap, and various other idiosyncrasies of JAX.

YouTube Video (Tutorial #1) <br/> Accompanying Jupyter Notebook <br/>

<p align="left"> <a href="https://www.youtube.com/watch?v=SstuvS-tVc0" target="_blank"><img src="https://img.youtube.com/vi/SstuvS-tVc0/0.jpg" alt="JAX from zero to hero!" width="480" height="360" border="10" /></a> </p>

Tutorial #2: From Hero to HeroPro+

In this video, we learn all additional components needed to train ML models (such as NNs) on multiple machines! We'll train a simple MLP model and we'll even train an ML model on 8 TPU cores!

YouTube Video (Tutorial #2) <br/> Accompanying Jupyter Notebook <br/>

<p align="left"> <a href="https://www.youtube.com/watch?v=CQQaifxuFcs" target="_blank"><img src="https://img.youtube.com/vi/CQQaifxuFcs/0.jpg" alt="JAX from Hero to HeroPro+!" width="480" height="360" border="10" /></a> </p>

Tutorial #3: Building a Neural Network from Scratch

Watch me code a Neural Network from scratch! :partying_face: In this 3rd video of the JAX tutorials series.

In this video, I build an MLP and train it as a classifier on MNIST using PyTorch's data loader (although it's trivial to use a more complex dataset) - all this in "pure" JAX (no Flax/Haiku/Optax).

I then do an additional analysis:

YouTube Video (Tutorial #3) <br/> Accompanying Jupyter Notebook (Note: I'll soon refactor it but I'll link the original)<br/>

<p align="left"> <a href="https://www.youtube.com/watch?v=6_PqUPxRmjY" target="_blank"><img src="https://img.youtube.com/vi/6_PqUPxRmjY/0.jpg" alt="Building a Neural Network from Scratch in pure JAX!" width="480" height="360" border="10" /></a> </p>

Tutorial #4: Machine Learning with Flax - From Zero to Hero

In this video, I cover everything you need to know to get started with Flax!

We cover init, apply, TrainState, etc. and other idiosyncrasies like the usage of mutable and rngs keywords.

YouTube Video (Tutorial #4) <br/> Accompanying Jupyter Notebook <br/>

<p align="left"> <a href="https://www.youtube.com/watch?v=5eUSmJvK8WA" target="_blank"><img src="https://img.youtube.com/vi/5eUSmJvK8WA/0.jpg" alt="Flax from Zero to Hero!" width="480" height="360" border="10" /></a> </p>

Tutorial #5 (coming up): Machine Learning with Haiku - From Zero to Hero

todo

Other useful content

Aside from the official docs here are some resources that helped me.

Videos

Blogs

Acknowledgements

Citation

If you find this content useful, please cite the following:

@misc{Gordic2021GetStartedWithJAX,
  author = {Gordić, Aleksa},
  title = {Get started with JAX},
  year = {2021},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/gordicaleksa/get-started-with-JAX}},
}

Connect With Me

If you'd love to have some more AI-related content in your life :nerd_face:, consider:

Licence

License: MIT