TPU (Tensor Processing Unit)
A Tensor Processing Unit (TPU) is a custom application-specific integrated circuit (ASIC) designed by Google to accelerate the matrix multiplications and tensor operations that dominate neural network training and inference. Unlike general-purpose GPUs, TPUs are built around a systolic array architecture that pipelines data through thousands of multiply-accumulate units with minimal memory-bandwidth overhead. Google has deployed multiple TPU generations internally and exposed them through Google Cloud, and they are the hardware on which models like Gemini and PaLM were trained.
Think of it like a specialized assembly line built to manufacture one product at enormous scale: a general factory (GPU) can produce many things but switches setups constantly, while the TPU line never reconfigures—it just runs the same multiply-accumulate operation billions of times per second.
# Training on TPU via JAX
import jax
print(jax.devices()) # [TpuDevice(id=0, ...), ...]
# JAX automatically maps array operations to the TPU's XLA compiler
PyTorch XLA and TensorFlow also support TPU backends, allowing most training code to run on TPUs with minimal changes.
TPUs occupy a specific niche in the hardware ecosystem alongside GPUs and NPUs; understanding them helps when choosing compute for Pretraining large Transformer models. They also influence Quantization choices, as TPU generations differ in their support for bfloat16 and int8 formats.
You'll encounter TPUs when using Google Cloud's AI infrastructure, working with JAX-based research code, or reading benchmark comparisons for Foundation Model training costs.