Unable to import jax in TPU

[paramjeet2021]

I am working on linux console and typing python takes me into the python console. When I use the following command in TPU machine

import jax

then it generates following mss and get out of the python prompt.

paramjeetsingh80@t1v-n-1c883486-w-0:~$ python3 Python 3.8.5 (default, Jan 27 2021, 15:41:15) [GCC 9.3.0] on linux Type "help", "copyright", "credits" or "license" for more information.

import jax 2021-07-08 17:41:39.660523: F external/org_tensorflow/tensorflow/core/tpu/tpu_executor_init_fns.inc:110] TpuTransferManager_ReadDynamicShapes not available in this library. Aborted (core dumped) paramjeetsingh80@t1v-n-1c883486-w-0:~$

This issue is causing problem in my code so I would like to figure out, what is this issue and how to get rid of this?

[jakevdp]

It might be that your system has the wrong version of libtpu. Try installing the version pinned here: https://github.com/google/jax/blob/38884b02f72950a8d187f81420a270e46afe8889/setup.py#L27

I believe the easiest way to install the correct version is to run

$ pip install -U jax
$ pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/jax_releases.html

[paramjeet2021]

Above command give some error but I researched and below command worked for me, but your answer give me the direction that it is a package issue.

pip install --upgrade pip

pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html