Mission 🚀

To provide 100 JAX exercises over different sections structured as a course or tutorials to teach and learn for beginners, intermediates as well as experts.

JAX

The JAX package in Python is a library for high performance and efficient machine learning research.

It is commonly used for various deep learning tasks and runs seamlessly on CPUs, GPUs as well as TPUs.

Exercises 📖

There are a total of 100 JAX exercises divided into 10 sets of Jupyter Notebooks with 10 exercises each. It is recommended to go through the exercises in order but you may start with any set depending on your expertise.

✅ Structured as exercises & tutorials - Choose your style
✅ Suitable for beginners, intermediates & experts - Choose your level
✅ Available on Colab, Kaggle, Binder & GitHub - Choose your platform
✅ Supports running on CPU, GPU & TPU - Choose your backend

All Sets • Exercises 1-100

Style	Colab	Kaggle	Binder	GitHub
Exercises
Solutions

Set 01 • JAX Introduction • Beginner • Exercises 1-10

Style	Colab	Kaggle	Binder	GitHub
Exercises
Solutions

Set 02 • Data Operations • Beginner • Exercises 11-20

Style	Colab	Kaggle	Binder	GitHub
Exercises
Solutions

Set 03 • Pseudorandom Numbers • Beginner • Exercises 21-30

Style	Colab	Kaggle	Binder	GitHub
Exercises
Solutions

Set 04 • Just-In-Time (JIT) Compilation • Beginner • Exercises 31-40

Style	Colab	Kaggle	Binder	GitHub
Exercises
Solutions

Set 05 • Control Flows • Beginner • Exercises 41-50

Style	Colab	Kaggle	Binder	GitHub
Exercises
Solutions

Set 06 • Automatic Differentiation • Intermediate • Exercises 51-60

Style	Colab	Kaggle	Binder	GitHub
Exercises
Solutions

Set 07 • Automatic Vectorization • Intermediate • Exercises 61-70

Set 08 • Pytrees • Intermediate • Exercises 71-80

Set 09 • Neural Networks • Expert • Exercises 81-90

Set 10 • Capstone Project • Expert • Exercises 91-100

The Jupyter Notebooks can also be run locally by cloning the repo and running on your local jupyter server.

git clone https://github.com/vopani/jaxton.git
python3 -m pip install notebook
jupyter notebook

P.S. The notebooks will be periodically updated to improve the exercises and support the latest version.

Contribution 🛠️

Please create an Issue for any improvements, suggestions or errors in the content.

You can also tag @vopani on Twitter for any other queries or feedback.

Credits 🙏

JAX

• Project
• Documentation

License 📋

This project is licensed under the Apache License 2.0.