Jax Examples for Beginners (Dive into ML Fundamentals)

A tutorial repository, JAX Examples for Beginners, offers notebooks and source code for learning JAX through worked examples, covering core transformations (grad, jit, vmap) and ecosystem libraries Flax (NNX API), Optax, Orbax and Grain. The examples range from basic operations and classic models — linear and logistic regression, nearest neighbor, K-Means, random forest, gradient boosted decision trees on the Boston Housing dataset, and Word2Vec trained on Wikipedia data — to neural networks built with Flax layers to classify MNIST digits, including a raw low-level implementation.

This tutorial was designed for easily diving into JAX https://docs.jax.dev/ , through examples. For readability, it includes both notebooks and source codes with explanation. It is suitable for beginners who want to find clear and concise examples about JAX. Besides the 'raw' JAX implementations jax.numpy , jax.grad , jax.jit , jax.vmap , ... , you can also find the most common libraries of the JAX ecosystem: - Flax https://flax.readthedocs.io/ NNX API : neural network layers and models nnx.Linear , nnx.Conv , nnx.LSTMCell , ... . - Optax https://optax.readthedocs.io/ : optimizers SGD, Adam, ... and loss functions. - Orbax https://orbax.readthedocs.io/ : save and restore models. - Grain https://google-grain.readthedocs.io/ : input data pipelines. - Basic Operations notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/1 Introduction/basic operations.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/1 Introduction/basic operations.py . A simple example that cover JAX basic operations. - JAX Transformations notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/1 Introduction/transformations.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/1 Introduction/transformations.py . Get started with JAX core transformations: grad automatic differentiation , jit compilation and vmap vectorization . - Linear Regression notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/2 BasicModels/linear regression.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/2 BasicModels/linear regression.py . Implement a Linear Regression with JAX. - Logistic Regression notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/2 BasicModels/logistic regression.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/2 BasicModels/logistic regression.py . Implement a Logistic Regression with JAX. - Nearest Neighbor notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/2 BasicModels/nearest neighbor.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/2 BasicModels/nearest neighbor.py . Implement Nearest Neighbor algorithm with JAX. - K-Means notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/2 BasicModels/kmeans.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/2 BasicModels/kmeans.py . Build a K-Means classifier with JAX. - Random Forest notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/2 BasicModels/random forest.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/2 BasicModels/random forest.py . Build a Random Forest classifier with JAX. - GBDT Gradient Boosted Decision Trees notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/2 BasicModels/gradient boosted trees.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/2 BasicModels/gradient boosted trees.py . Implement a Gradient Boosted Decision Trees with JAX to predict house value using Boston Housing dataset. - Word2Vec Word Embedding notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/2 BasicModels/word2vec.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/2 BasicModels/word2vec.py . Build a Word Embedding Model Word2Vec from Wikipedia data, with JAX. - Simple Neural Network notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/3 NeuralNetworks/neural network.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/3 NeuralNetworks/neural network.py . Use Flax 'layers' and 'model' API to build a simple neural network to classify MNIST digits dataset. - Simple Neural Network low-level notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/3 NeuralNetworks/neural network raw.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/3 NeuralNetworks/neural network raw.py . Raw implementation of a simple neural network to classify MNIST digits dataset. - Convolutional Neural Network notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/3 NeuralNetworks/convolutional network.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/3 NeuralNetworks/convolutional network.py . Use Flax 'layers' and 'model' API to build a convolutional neural network to classify MNIST digits dataset. - Convolutional Neural Network low-level notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/3 NeuralNetworks/convolutional network raw.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/3 NeuralNetworks/convolutional network raw.py . Raw implementation of a convolutional neural network to classify MNIST digits dataset. - Recurrent Neural Network LSTM notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/3 NeuralNetworks/recurrent network.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/3 NeuralNetworks/recurrent network.py . Build a recurrent neural network LSTM to classify MNIST digits dataset, using Flax 'layers' and 'model' API. - Bi-directional Recurrent Neural Network LSTM notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/3 NeuralNetworks/bidirectional rnn.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/3 NeuralNetworks/bidirectional rnn.py . Build a bi-directional recurrent neural network LSTM to classify MNIST digits dataset, using Flax 'layers' and 'model' API. - Dynamic Recurrent Neural Network LSTM notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/3 NeuralNetworks/dynamic rnn.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/3 NeuralNetworks/dynamic rnn.py . Build a recurrent neural network LSTM that performs dynamic calculation to classify sequences of variable length, using Flax 'layers' and 'model' API. - Auto-Encoder notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/3 NeuralNetworks/autoencoder.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/3 NeuralNetworks/autoencoder.py . Build an auto-encoder to encode an image to a lower dimension and re-construct it. - Variational Auto-Encoder notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/3 NeuralNetworks/variational autoencoder.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/3 NeuralNetworks/variational autoencoder.py . Build a variational auto-encoder VAE , to encode and generate images from noise. - GAN Generative Adversarial Networks notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/3 NeuralNetworks/gan.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/3 NeuralNetworks/gan.py . Build a Generative Adversarial Network GAN to generate images from noise. - DCGAN Deep Convolutional Generative Adversarial Networks notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/3 NeuralNetworks/dcgan.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/3 NeuralNetworks/dcgan.py . Build a Deep Convolutional Generative Adversarial Network DCGAN to generate images from noise. - Save and Restore a model notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/4 Utils/save restore model.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/4 Utils/save restore model.py . Save and Restore a model with JAX, using Orbax. - Build Custom Layers & Modules notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/4 Utils/build custom layers.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/4 Utils/build custom layers.py . Learn how to build your own layers / modules and integrate them into Flax Models. - Tensorboard notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/4 Utils/tensorboard.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/4 Utils/tensorboard.py . Track and visualize neural network computation graph, metrics, weights and more using TensorBoard. - Load and Parse data notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/5 DataManagement/load data.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/5 DataManagement/load data.py . Build efficient data pipeline with Grain Numpy arrays, Images, CSV files, custom data, ... . - Build an image dataset notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/5 DataManagement/build an image dataset.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/5 DataManagement/build an image dataset.py . Build your own images dataset from image folders, and train a model on it. - Build and Load ArrayRecords notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/5 DataManagement/array records.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/5 DataManagement/array records.py . Convert data into ArrayRecord format the record format of Grain , and load them. - Image Transformation i.e. Image Augmentation notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/5 DataManagement/image transformation.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/5 DataManagement/image transformation.py . Apply various image augmentation techniques with JAX, to generate distorted images for training. - Multi-Device Basics notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/6 Hardware/multi device basics.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/6 Hardware/multi device basics.py . A simple example to introduce computation on multiple devices GPUs, TPUs with JAX. - Multi-Device Training notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/6 Hardware/multi device training.ipynb code https://github.com/aymericdamien/JAX-Examples/blob/master/6 Hardware/multi device training.py . Train a convolutional neural network with multiple devices GPUs, TPUs on CIFAR-10 dataset. Some examples require MNIST dataset for training and testing. Don't worry, this dataset will automatically be downloaded when running examples with the small helper common/datasets.py https://github.com/aymericdamien/JAX-Examples/blob/master/common/datasets.py , files are cached in a data/ folder . MNIST is a database of handwritten digits, for a quick description of that dataset, you can check this notebook https://github.com/aymericdamien/JAX-Examples/blob/master/notebooks/0 Prerequisite/mnist dataset intro.ipynb . Official Website: http://yann.lecun.com/exdb/mnist/ http://yann.lecun.com/exdb/mnist/ . To download all the examples, simply clone this repository: git clone https://github.com/aymericdamien/JAX-Examples To run them, you also need JAX and a few libraries of its ecosystem. To install them: pip install -r requirements.txt or with NVIDIA GPU support : pip install -r requirements.txt "jax cuda13 " For older NVIDIA GPUs or drivers e.g. V100 GPUs, or drivers older than 580 , use "jax cuda12 " instead. For more details about JAX installation GPU, TPU, ... , you can check JAX Installation Guide https://docs.jax.dev/en/latest/installation.html . This is the JAX version of TensorFlow-Examples https://github.com/aymericdamien/TensorFlow-Examples . Every example follows the same structure, so you can also compare both side by side.