Porting nanochat to a TPU: what carries over from PyTorch, and what breaks

Karpathy's nanochat normally runs on an 8×H100 GPU node, and several ports of it to JAX already exist. Among them, my aim was to keep the config and architecture as close to nanochat as possible (p...