I Rebuilt Karpathy's NanoChat in JAX. Here's What XLA Gets Right and What It Gets Dead Wrong.

AI GDE TPU Sprint 2026 · Google TPU Research Cloud Quick summary: We ported Andrej Karpathy's NanoChat architecture from PyTorch to JAX and Flax NNX. The repo is about 12,400 lines across source and scripts. We trained a nano model (885K parameters) on TinyStories in under 10 minutes on a single GPU and served it through a streaming chat UI. XLA compilation eliminates Python overhead after a one-time upfront cost. The same code runs on TPU without modification. The catch: no vLLM, no Flash Atten