specify jax default device to use

jax_default_device.py

import jax
import jax.numpy as jnp

jax.config.update("jax_default_device", jax.devices("cpu")[0])

#runs on cpu
long_vector = jnp.arange(int(1e7))
%timeit jnp.dot(long_vector, long_vector).block_until_ready()

jax.config.update("jax_default_device", jax.devices("gpu")[2])

#runs on gpu
long_vector = jnp.arange(int(1e7))
%timeit jnp.dot(long_vector, long_vector).block_until_ready()