Jaxformer
Preface
I love Richard Feyman and his lectures are a big influence for why I studied mathematics (because I didn’t understand it). I think this quote basically sums up why I wrote this article.
“If you can’t explain something to a first year student, then you haven’t really understood.”
I want to push myself to understand these AI models as deeply as I possibly can, because it seems likely they are the most sigificant technological development in history, so it will probably come in handy.
A point on AI assistance
I am a strong believer in using AI to assist with learning, coding, and reasearching. However, AI prose is extremely noticable and unenjoyable to read. I have elected to not have Claude write a single word of this series and in the jaxformer repository I didn’t use Claude to write any of the model code, I did use Claude to make getting everything to run in google colab easier. I have used Claude to review my writing, to make the diagrams you see.
Further reading
What I would recommend reading afterwards or if you want to go deeper:
- https://jax-ml.github.io/scaling-book/
- https://rlvrbook.com/
- https://rlhfbook.com/
- https://djdumpling.github.io/2026/01/31/frontier_training.html