Does JAX come with XLA binaries that I can use?

[hayden-donnelly]

I want to do some stuff with XLA in C++, but it's taking a very long time to build XLA from source. To get around this I figured I might be able to use whichever XLA binaries that my version of JAX is already using, is this possible?

For context, I've installed JAX with the following command:

python3 -m pip install "jax[cuda11_cudnn86]" -f  "https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"

And my /usr/local/lib/python3.8/dist-packages/jaxlib directory looks like this:

__init__.py
__pycache__
cpu
cpu_feature_guard.so
cuda
ducc_fft.py
gpu_linalg.py
gpu_prng.py
gpu_rnn.py
gpu_solver.py
gpu_sparse.py
gpu_triton.py
hlo_helpers.py
lapack.py
mlir
py.typed
utils.so
version.py
xla_client.py
xla_extension
xla_extension.so

I think xla_extension.so may be what I'm looking for, but I'm not sure.

[hawkinsp]

What you want is buried inside xla_extension.so, but not with an interface that is useful to you most likely (since that module exposes a Python API).

There is a better way, which is that the XLA-PJRT runtime that JAX uses has a C API: https://github.com/openxla/xla/blob/main/xla/pjrt/c/README.md

Some JAX plugins already use this API (Google TPU, Apple Metal, to name two). We are planning to shortly update JAX's NVIDIA GPU plugin to also use this API, and at that point JAX will have a build of that plugin API as a .so file.

[hayden-donnelly]

Awesome, I'll check it out. Thank you.