Cloud TPU provides Google Cloud virtual machines running directly on hosts attached to TPU accelerators. It suits large neural-network training and lets JAX move beyond local CPU or GPU execution. section 3.2.5
A practical Cloud TPU workflow includes:
- Create and connect to a TPU VM with gcloud, then expose its Jupyter server through a Colab local runtime. section 3.2.5
- Install the JAX TPU package and inspect the selected backend and local TPU devices. section 3.2.5
- Delete the VM afterward to avoid ongoing Google Cloud charges. section 3.2.5
$gcloud compute tpus tpu-vm create node-jax \ --zone us-central1-b --accelerator-type v2-8 \ --version tpu-vm-base $gcloud compute tpus tpu-vm list --zone us-central1-b $gcloud compute tpus tpu-vm ssh \ --zone us-central1-b node-jax -- \ -L 8888:localhost:8888
Use the gcloud commands to create, list, and SSH into a TPU VM before connecting a notebook runtime. section 3.2.5
Check the JAX TPU backend
Install the TPU-enabled JAX package, print the active platform, and call jax.local_devices() to confirm that TPU devices are available. section 3.2.5
