TPUs are amazing, but in my experience, debugging issues with them can be a bit tricky. Since nvidia's gpus are more common place (especially outside gcp), you can find a lot more information when you get stuck, it's also more battle tested, etc.
Of course that assumes your issues aren't with the jit itself or inside pmap, etc. That shit's hard.