如何只安装 XLA？

Posted 2023-02-16

【中文标题】如何只安装 XLA？

【英文标题】：How to install just XLA?

【发布时间】：2021-09-18 05:47:41

【问题描述】：

我想使用 XLA 作为我项目的后端。有没有推荐的方法来独立安装它（没有 TensorFlow 的其余部分）。 Jax 可能会这样做，但查看他们的存储库并不清楚是如何做到的。

更新我为此使用 TensorFlow 提出了a ticket

【参考方案1】：

除了 tensorflow 之外，没有支持自行安装 XLA 的方法。

也就是说，JAX 确实在 jaxlib 包中与 tensorflow 分开提取、构建和捆绑 XLA。您可以在此处查看各种平台上 jaxlib 的相关构建脚本：https://github.com/google/jax/tree/main/build

特别是查看build_wheel.py，其中包含作为 jaxlib 构建的一部分从 tensorflow 源中提取 XLA 相关片段的脚本。

【讨论】：

谢谢。这是一个相当不错的构建脚本。现在我已经用 TensorFlow 提出了 a ticket 让 XLA 可以自己使用

啊，我明白了 - 你的意思是 jaxlib 本质上 is 是一个独立的 XLA 吗？还是 Python 包装的 XLA？我猜是后者，因为它是 pip 可安装的。我正在寻找理想的 C/C++ API

不完全是 - jaxlib 并非设计为 XLA 的通用版本。也就是说，在某种程度上可以通过 C++ 与 XLA API 进行交互：例如，参见 dfm.io/posts/extending-jax

【参考方案2】：

elixir-nx 项目已发布XLA binaries。

TensorFlow 项目的其他部分提供了很多功能，可能不是 TensorFlow 特有的，但与它并存。 IIUC，XLA接口在tensorflow/compiler/xla。