IMPORTANT UPDATE (May 26, 2021): Cloud TPU VMs are now publicly available and you do not need to fill out this form to get access!
We're glad you're interested in trying JAX with the new Cloud TPU experience! This is still a "preview" release, meaning there may still be some rough edges, but we hope you’ll give JAX a try on the new Cloud TPU experience and let us know what you think.
To get started, check out the Cloud TPU Preview documentation here:
https://cloud.google.com/tpu/docs/users-guide-tpu-vm. We recommend the JAX quickstart in particular:
https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm.
If you need to request additional Cloud TPU quota, more information is available here
https://cloud.google.com/tpu/docs/quota. If applicable, you may also be interested in signing up for the TPU Research Cloud program to receive free Cloud TPU quota:
https://sites.research.google/trc/.
If you run into any issues that are specific to JAX, please file an issue (
https://github.com/google/jax/issues) or post a discussion question (
https://github.com/google/jax/discussions). You’re also welcome to email us directly at
jax-cloud-tpu-team@google.com.
If you run into general Cloud TPU issues (e.g. trouble creating a Cloud TPU VM), or you’re not sure if the issue is specific to JAX, please email
cloud-tpu-support@google.com, or
trc-support@google.com if you are a TRC member.
We’d also love to hear any feedback you have, good or bad, at
jax-cloud-tpu-team@google.com. Hope to hear from you!