无法在 Google TPU 中导入 python 包 jax答案

【问题标题】：Not able to import python package jax in Google TPU无法在 Google TPU 中导入 python 包 jax
【发布时间】：2021-07-08 17:49:37
【问题描述】：

我正在使用 linux 控制台，输入 python 会将我带入 python 控制台。当我在 TPU 机器中使用以下命令时

import jax

然后它会生成以下 mss 并退出 python 提示符。

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:~$

这个问题导致我的代码出现问题，所以我想弄清楚这是什么问题以及如何解决这个问题？

【问题讨论】：

标签： python-3.x tpu jax spacy-transformers sentence-transformers

【解决方案1】：

可能是您的系统没有正确版本的 libtpu。尝试安装列出的版本here。

您应该能够自动执行此操作

$ pip install -U pip  # older pip may not support extra requirements
$ pip install -U jax  # newer jax required for [tpu] extras declaration
$ pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/jax_releases.html

【讨论】：

libtpu 版本随时间而变化。当您更新 jax、pytorch/xla 等时，您需要检查 libtpu 版本是否保持一致。

【解决方案2】：

上面的命令给出了一些错误，但我研究过，下面的命令对我有用。但你的回答给了我方向，这是一个包裹问题。

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

【讨论】：

这应该是对其他答案的评论