It took roughly 2-3 weeks to translate Torch to JAX, but I had past experience writing JAX from my time at Google.

We spent nearly 4 weeks getting PyTorch XLA working on TPU. Hope that answers your question!