Finetuning Wav2Vec2 on TPU

I am planning to finetune a Wav2Vec2-XLSR-53 on a newly generated speech dataset. Are there any example scripts on how to finetune Wav2Vec2 on TPUs? I do find some PyTorch examples but these are for GPUs only.

I would prefer using Flax for this but the only scripts I am able to find is for pretraining.