XLA is the compatibility layer. You train using Torch/Jax/whatever you want, and XLA ensures you can use TPU. There aren't many hurdles, most people train in the cloud anyway.