How to install just XLA?

Viewed 214

I want to use XLA as a backend for my project. Is there a recommended way to install it on its own (without the rest of TensorFlow). Jax probably does this, but looking in their repository it's not obvious how.

UPDATE I raised a ticket with TensorFlow for this

2 Answers

There is no supported way to install XLA on its own apart from tensorflow.

That said, JAX does extract, build, and bundle XLA separately from tensorflow within the jaxlib package. You can see the relevant build scripts for jaxlib on various platforms here: https://github.com/google/jax/tree/main/build

In particular, take a look at build_wheel.py, which contains the scripts that extract relevant pieces of XLA from the tensorflow source as part of the jaxlib build.

The elixir-nx project has released XLA binaries.

The XLA interface is in tensorflow/compiler/xla. The rest, IIUC, is supporting functionality.

Related