Are provided benchmarks CPU only

[rahulbatra85]

There are some benchmarks available here https://github.com/google/jax/tree/main/benchmarks
Are they meant to be CPU only?

[jakevdp]

Thanks for the question! We use these same benchmarks on CPU, GPU, and TPU.

[rahulbatra85]

Thanks for the response!
When I run these benchmarks they only use CPU. I want to target a different backend(GPU). Is their some setup needed to target a different backend? or should all that be handled by google_benchmark module? I would appreciate if you can share some details on this. Thanks!

[jakevdp]

If you install a GPU-enabled jaxlib, then the benchmarks should run on GPU. See https://github.com/google/jax#pip-installation-gpu-cuda for GPU installation instructions.