JAX is not exclusive to TPUs, but it shines on that hardware because each TPU server has 8 TPU accelerators working in parallel.

Setup
	
First make sure diffusers is installed.


	Copied
!pip install jax==0.3.25 jaxlib==0.3.25 flax transformers ftfy
!pip install diffusers


	Copied
import jax.tools.colab_tpu

jax.tools.colab_tpu.setup_tpu()
import jax


	Copied
num_devices = jax.device_count()
device_type = jax.devices()[0].device_kind

print(