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 Attention 3, and painful debugging inside JIT-compiled functions. This post covers what worked, what did not, and when you should care.

The nanochat-jax transformer block.…