Microjax: JAX in two classes and six functions

or Read on Github (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.