If you are running this guide in Colab, select *Runtime* in the menu above, select the option *Change runtime type*, and then select *TPU* under the *Hardware accelerator* setting. Import JAX and quickly check whether you're using a TPU:  
```python
import jax
import jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu()