Why is the TPU not recognized on my Google Cloud TPU VM instance?

Question

I have launched a Google Cloud TPU VM instance and installed the latest version of JAX, but it cannot see my TPU. Following the instructions at https://cloud.google.com/tpu/docs/troubleshooting/trouble-jax I encounter the following:

>>> import jax
>>> jax.devices()
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
[CpuDevice(id=0)]
>>> TF_CPP_MIN_LOG_LEVEL=0
>>> jax.devices()
[CpuDevice(id=0)]

All of the Google Search results I have seen for this error suggest installing JAX with CUDA support, but shouldn't that be unnecessary with TPUs?

Could you share what TPU VM version are you using and what commands you ran to reproduce this? — Allen Wang, Jul 26 '22 at 19:51
can you run `TF_CPP_MIN_LOG_LEVEL=0 python3` and then import jax, so that we can see better logs there? — Shiva Shahrokhi, Aug 01 '22 at 18:59
Am having the same problem, am getting `TPU platform initialization failed: NOT_FOUND: No ba16c7433 device found.` — Chris Flesher, Sep 19 '22 at 11:13

Chris Flesher · Answer 1 · 2022-11-19T00:43:20.863

1

I recommend to upgrade the jax version.

pip3 install -u jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

EDIT: Actually it seems like this is a bug, refer to:

https://github.com/google/jax/issues/13260

edited Nov 19 '22 at 00:43

answered Nov 15 '22 at 17:24

Chris Flesher

987
1
10
13

Why is the TPU not recognized on my Google Cloud TPU VM instance?

1 Answers1