Get started with JAX by Aleksa Gordić
A collection of tutorials and resources to help developers learn JAX, Flax, and Haiku for machine learning.
What is Get started with JAX by Aleksa Gordić?
Get Started with JAX is a learning resource repository designed to help developers master the JAX ecosystem, including Flax and Haiku. It provides a series of video tutorials and Jupyter notebooks that cover everything from basic JAX operations to building neural networks from scratch and training on multiple TPU cores. The project aims to lower the entry barrier for using JAX as an alternative to established frameworks like PyTorch and TensorFlow.
Machine learning practitioners, researchers, and students who want to learn JAX, Flax, and Haiku for developing and training neural networks. It's especially useful for those transitioning from PyTorch or TensorFlow.
It offers a curated, hands-on learning path with practical examples and Colab-ready notebooks, saving learners time by focusing on high-quality, personally vetted content rather than overwhelming them with every available resource.
Overview
The purpose of this repo is to make it easy to get started with JAX, Flax, and Haiku. It contains my "Machine Learning with JAX" series of tutorials (YouTube videos and Jupyter Notebooks) as well as the content I found useful while learning about the JAX ecosystem.
Use Cases
Best For
Related Projects
Found a gem we're missing?
Open-Awesome is built by the community, for the community. Submit a project, suggest an awesome list, or help improve the catalog on GitHub.
