Not sure if this is needed, but in the terminal session you are about to run your training script in, you might want to type:
export XRT_TPU_CONFIG="localservice;0;localhost:51011"



AFAIK, this is needed if you are using pytorch/xla, and not needed for flax/jax.
You might need to do export USE_TORCH=False if you have torch installed but want to use flax.