JAX-style automatic-differentiation system from first principles in fewer than 500 lines of Python. It covers computation tracing, NumPy-like primitives, vector-Jacobian products, reverse-mode autodiff, and higher-order differentiation, with a small neural-network library demonstrating end-to-end training.
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.