If you have access to a TPU, the Flax training script runs even faster! Let’s run the training script on the Google Cloud TPU VM. Create a single TPU v4-8 VM and connect to it:   Copied ZONE=us-central2-b