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

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!