Scaling LLMs by Hand ✍️ in JAX (Aug 13, 2025)
AI by Hand ✍️ Seminars
When Google’s JAX team first approached me to sponsor a special lecture, I told them, “You asked the wrong person. I’ve never written a single line of JAX code, but I can walk people through the math and algorithms behind the syntax. Is that okay?”
They said, “Absolutely. That’s exactly what we want you to do.”
In preparing for this lecture, I got to meet many people on Google’s JAX team, including the creators. At one JAX event, Roy Frostig, one of JAX’s creators, opened his keynote with a single slide: “Matrix Multiplication is Everything!”
If you know me, you know I couldn’t agree more.
In this seminar, I walk you through:
How attention scales to millions of tokens
What vectorization really means in JAX
How to understand @vmap (vectorization map)
All by hand ✍️ in Excel, so you can see exactly what’s going on inside the black box.
P.S. This special lecture is sponsored by Google. If you’d like to sponsor a future AI by Hand ✍️ lecture to teach an important topic to a large, global AI audience, feel free to reach out.
Feedback
Recording
Become a paid subscriber to access the full recording.