jax can connect to tpu v3-8, but not v3-32 and v3-64

[lucasliunju]

jax can connect to tpu v3-8, but not v3-32 and v3-128.

• [ jax connect to tpu, I try to connect tpu follow this code:]

import requests
import os
url = f'http://[ip address of tpu]:8475/requestversion/tpu_driver_nightly'
requests.post(url)
TPU_DRIVER_MODE = 1

from jax.config import config
# The following is required to use TPU Driver as JAX's backend.
config.FLAGS.jax_xla_backend = "tpu_driver"
config.FLAGS.jax_backend_target = "grpc://[ip address]:8470"

• [ RuntimeError: Deadline exceeded: Failed to connect to remote server at address: grpc://10.59.106.169:8470. Error fr
  om gRPC: Deadline Exceeded. Details: ]

[lucasliunju]

Hi @mattjj and @skye,

As shown in the above, I try to connect jax and tpu. It works when I use v3-8 version of tpu. However, when I try to connect jax and multi-host tpu (such as v3-128) I find it cannot work when I try to :

print(jax.devices())

The error is :

RuntimeError: Deadline exceeded: Failed to connect to remote server at address: grpc://10.59.106.169:8470. Error fr
om gRPC: Deadline Exceeded. Details:

Currently, I find many great work (vision transformer, sam) is based on jax and multi-host tpu. Therefore, I would like to reproduce these work and learn to use jax. During this process, I think jax is very intertsting and will become more popular and gain more attention.

Thank you very much!

I am looking forward to your reply!

Best regards,
Yong