From bigrams to GPT-2, one component at a time (in Jax) (gilesthomas.com)