For anybody else trying to grok backprop, I found it very helpful to look at karpathy's microgpt, a barebones LLM in 200 lines of python. The autograd itself is just 40 lines, and reading that code and then trying to do a clean room implementation of it is a great exercise: