Comparing Jax on TPU vs Numpy on CPU does not make any sense. Of course a GEMM hardware accelerator will be much faster than a general purpose CPU. What is the

Comparing Jax on TPU vs Numpy on CPU does not make any sense. Of course a GEMM hardware accelerator will be much faster than a general purpose CPU. What is the point of this comparison? You either run the same code on different hardware, or different code on the same hardware. A much more interesting comparison would be CuPy or Pytorch code vs Jax code running on A100.