Anything that you can express as 128x128 (but ideally much larger) dense matrix multiplication and nothing else

I think https://jax-ml.github.io/scaling-book/ is one of the best references to go through. It details how single device and distributed computations map to TPU hardware features. The emphasis is on mapping the transformer computations, both forwards and backwards, so requires some familiarity with how transformer networks are structured.

definitely closed-loop, might even be a servo

From the acknowledgment at the end, I guess the author has access to TPUs through https://sites.research.google/trc/about/

This is not the only way though. TPUs are available to companies operating on GCP as an alternative to GPUs with a different price/performance point. That is another way to get hands-on experience with TPUs.

Everything thats in the blog post is basically well known already. Google publishes papers and gives talks about their TPUs. Many details are lacking though, and require some assumptions/best guesses. Jax and XLA are (partially) open source and give clues about how TPUs work under the hood as well.

https://arxiv.org/abs/2304.01433

https://jax-ml.github.io/scaling-book/