Amazing news, I remember the days when you had to fight XLA to make PyTorch run on TPUs, it seems we finally got first-party support!

Since HW specifics live mostly in the kernel layer (Pallas or CUDA), I imagine the Jax vs PyTorch question is more of a style preference now?