Special Lecture: Scaling LLMs in JAX (sponsored by Google)
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.
Next Wednesday, I’ll 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.
900+ people have already registered!
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.