Gemma 4 Runs in Rust on a TPU via TorchTPU, Trailing Python Speed but Using Less Memory

A guide demonstrates running the Gemma 4 E2B model with its code written in Rust on a single TPU v5e chip in Google Colab, with every operation executed by TorchTPU and no JAX involved. A Rust crate handles checkpoint loading, 4-bit weight unpacking, the 35 layers, the ring cache, and the decode loop, while Python hosts the backend and starts the engine. In benchmarks on the same chip and prompt, the Rust versions produced the same tokens as compiled PyTorch but streamed at roughly half the speed while using less than half the peak memory. Because TorchTPU provides no supported C++ API, Python must host the process, and a pyo3 extension bridges Rust calls to the same libtorch dispatcher. The author notes that the tch crate has no TPU device type, so tensors are moved to the chip by matching an existing tensor's device.