This repo contains code for the Nova Science coding group.
This online book is quickly becoming the standard text for Deep Learning.
And here is repo with list of projects in JAX.
The mnist.ipynb notebook builds a simple MNIST classifier. Here is a Google Colab version. And for a classification over 99.5% please consult the Jiucheng branch.
The vae_mnist.ipynb notebook builds a very simple Variational Autoencoder for the MNIST dataset. Here is a Colab version. And here is a very readable elementary introduction to VAE.
The ppo.ipynb notebook trains an agent to control a cart-pole with reinforcement learning. Here is a Colab version. This blog has a simple introduction to the PPO algorithm. The notebook is an adaptation of a notebook from this blog which also provides evidence for the extereme speedups and new opportunities afforded by using JAX end-to-end in reinforcement learning.
The diffusion.ipynb notebook trains a basic Denoising Diffusion probabilistic model.
The rng.ipynb notebook explains Pseudo Random Number Generation in JAX. Here is a Google Colab version.
The autodiff.ipynb looks of examples of how autodiff is used in JAX. Here is a Google Colab version.
The pytrees.ipynb looks at structure and utilities for pytrees. Here is a Google Colab version.
The state.ipynb notebook investigates how program state is handled in purely functional JAX. Here is the Colab version.
The jtransform.ipynb notebook investigates how JAX tranforms functions based on their intermediate langage jaxpr representations. Here is the Colab version.