Machine Learning with JAX Tutorials

Get started with JAX! 💻 ⚡

The goal of this repo is to make it easier to get started with JAX!

JAX is becoming an increasingly popular alternative to PyTorch and TensorFlow. 😎

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.

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

YouTube Video.
Accompanying Jupyter Notebook.

Other useful content

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

Videos

• Introduction to JAX (gives a very high-level overview)
• JAX: Accelerated Machine Learning Research | SciPy 2020 | VanderPlas (many more details)
• NeurIPS 2020: JAX Ecosystem Meetup (DeepMind team about the ecosystem of libs around JAX)
• Introduction to JAX for Machine Learning and More (nice, hands-on workshop)
• Day 1 Talks: JAX, Flax & Transformers | HuggingFace (all 4 talks are good)
• Day 2 Talks: JAX, Flax & Transformers | HuggingFace (only the first 2 talks)

Blogs

• Using JAX to accelerate our research | DeepMind (similar info as the NeuroIPS 2020 video)
• You don't know JAX | Colin Raffel

Acknowledgements

• The notebooks were heavily inspired by the official JAX docs.

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 🤓 , consider:

• Subscribing to my YouTube channel The AI Epiphany 🔔
• Follow me on LinkedIn and Twitter 💡
• Follow me on Medium 📚 ❤️
• Join the Discord community! 👪

Licence