If you read the jax source code you'll hit something called xla_client. Often imported like this
from . import xla_client
This implies that xla_client is a python module, but I can't find any file with that name or reference to a variable of that name.
I assume that it is related to https://pypi.org/project/jaxlib/, but this package just links back to the jax source code.
Can anybody clue me in?
The file you're referring to is stored at https://github.com/tensorflow/tensorflow/tree/master/tensorflow/compiler/xla/python
Let me expound further: xla_client is partly a wrapper around a specially compiled c++ file called xla_extension.so, for example see
from . import xla_extension as _xla
and numerous references to _xla throughout xla_config. The source for this file is https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/python/xla.cc, which we know because it says so quite clearly in https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/python/BUILD
pybind_extension(
name = "xla_extension",
srcs = [
"xla.cc",
],
...