The JAX docs have a good explanation for how a TPU works

https://docs.jax.dev/en/latest/pallas/tpu/details.html#what-...

Its not really useful for other workloads (unless your workload looks like a bunch of matrix multiplications).