jax 库中的“xla_client”到底是什么？

Question

如果您阅读jax source code，您会看到名为xla_client 的内容。经常这样导入

from . import xla_client

这意味着xla_client 是一个python 模块，但我找不到任何具有该名称的文件或对该名称的变量的引用。

我假设它与https://pypi.org/project/jaxlib/ 相关，但这个包只是链接回 jax 源代码。

谁能帮我解答一下？

Answer 1

您引用的文件存储在https://github.com/tensorflow/tensorflow/tree/master/tensorflow/compiler/xla/python

让我进一步解释一下：xla_client 部分是对名为xla_extension.so 的专门编译的 c++ 文件的包装，例如，请参阅

from . import xla_extension as _xla

在整个xla_config 中多次引用_xla。这个文件的来源是https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/python/xla.cc，我们知道是因为它在https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/python/BUILD中说得很清楚

pybind_extension(
name = "xla_extension",
srcs = [
    "xla.cc",
],
...