Comment by drdirk

Jax uses the XLA compiler which is compatible with GPU and CPU.