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

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