Porting NanoChat to JAX: What XLA Excels At—and Where It Fails

This deep-dive port of NanoChat from PyTorch to JAX exposes both the promise and pitfalls of XLA compilation for AI development. While XLA enables GPU/TPU code portability and eliminates Python overhead, it comes with significant trade-offs—limited library support, painful debugging inside JIT functions, and the absence of critical optimizations like Flash Attention 3. The project offers concrete lessons on when JAX's universal compilation strategy pays off and when it doesn't.