tensorflow - 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",
],
...