This article is about running a hand-written Gemma 4 port in pure JAX on three different accelerators, and about the two places the abstraction leaks.

The code is here:

github.com/xbill9/gemma4-dev

What is this project trying to Do?

This project aims to serve one Gemma 4 checkpoint from one JAX port across every accelerator I can rent, and to find out — by measurement, not by reading docs — which parts of "it's just JAX" are true.