It depends on the workload. Google's own guidance says TPUs suit models "dominated by matrix computations", long training runs and large models with large batch sizes, and that GPUs suit models with many custom PyTorch or JAX operations.
What Google's comparison says
From Google's TPU introduction:
| Hardware | Suited to |
|---|---|
| CPU | Quick prototyping, simple and small models |
| GPU | Models with many custom PyTorch/JAX operations, medium to large models |
| TPU | Matrix-dominated models, long training runs, large batch sizes |
Google adds that TPUs are not suited to workloads that need high-precision arithmetic or frequent branching.
The practical trade-offs
TPU code goes through the XLA compiler, and TPUs are a Google service, covered in how to access a TPU. On the software side, PyTorch runs through TorchTPU.
For the architecture and cost view, read TPU vs GPU, and for the whole field, NVIDIA's main competitors in AI. To try a workload on a GPU, see GPU pricing and availability.
