Microjax: JAX in two classes and six functions

<a target="_blank" href="https://colab.research.google.com/github/joelburget/microjax/blob/main/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a> or <a href="https://github.com/joelburget/microjax/blob/main/tutorial.ipynb">Read on Github</a> (I recommend actually running the notebook, either on your own computer or Colab).

This is inspired by Andrej Karpathy's Micrograd, a PyTorch-like library in about 150 lines of code. Despite PyTorch's popularity, I prefer the way JAX works because it a more functional style.

This tutorial borrows heavily from Matthew J Johnson's great 2017 presentation on the predecessor to JAX, autograd: Video / Slides / Code. My main contribution is simplifying a bit and packaging it as a notebook.