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

Question

如果您阅读jax 源代码，您会遇到名为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部分是一个名为 的专门编译的 C++ 文件的包装器xla_extension.so，例如参见

from . import xla_extension as _xla

_xla以及全文的大量引用xla_config。该文件的来源是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",
],
...