如何只安装 XLA?

How to install just XLA?

我想使用 XLA 作为我项目的后端。是否有推荐的方法自行安装(没有其余的 TensorFlow)。 Jax 可能会这样做,但查看他们的存储库并不清楚如何做。

更新 我为此

使用 TensorFlow 提出 a ticket

除了 tensorflow 之外,不支持单独安装 XLA。

也就是说,JAX 确实在 jaxlib package. You can see the relevant build scripts for jaxlib on various platforms here: https://github.com/google/jax/tree/main/build

中与 tensorflow 分开提取、构建和捆绑 XLA

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

elixir-nx project has released XLA binaries.

其中有很多来自 TensorFlow 项目其余部分的功能,可能不是 TensorFlow 特有但与它并存的功能。 IIUC,XLA界面在tensorflow/compiler/xla.