Jax: Cloud TPU Support

[aidangomez]

Hello all,

Super stoked about this project and so glad it's out in the open now! I just wanted to make an issue people can follow to track the support of Cloud TPUs in JAX.

As soon as this is ready to be tested for.ai and I would be super eager to give it a try and help out!

All the best,
Aidan

[mattjj]

Great idea! Thanks for opening this. We're working on Cloud TPU support, and will use this issue to track updates.

[Superman132]

Does cloud tpu support mean colaboratory tpu support, or is that a seperate issue?

[mattjj]

They're related, and will probably be solved simultaneously.

[dynamicwebpaige]

Pinging for a status check / update. 🙂

[joaogui1]

Guess we can close this? @mattjj, @skye

[skye]

Thanks for the ping @joaogui1 !

To everyone following this thread, JAX can now be used with Cloud TPUs! It's still early in development, so we'd appreciate your help trying it out. See https://github.com/google/jax/tree/master/cloud_tpu_colabs.