gcloud alpha compute tpus tpu-vm ssh $VM_NAME --zone $ZONE -- \
```  
Install JAX 0.4.5:  
```bash
pip install "jax[tpu]==0.4.5" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
```  
Then install the required dependencies for the Flax script:  
```bash
cd examples/controlnet
pip install -r requirements_flax.txt
```  
</hfoption>
</hfoptions>  
<Tip>