How to compile JAX on a TPU instance

[awav]

Hello everyone,

I'm trying to compile JAX with a custom XLA on the TPU VM, but I'm getting an error showed below. I would appreciate it if someone could help me to figure out what is the correct procedure for compiling JAX on TPUs.

I'm building JAX with this command (python with anaconda and python3 without):

> cd ~/jax
> python build/build.py --enable_tpu
> pip install --force-reinstall dist/jaxlib-*.whl

The import error after successfull compilation:

In [1]: import jax
2022-01-17 20:08:33.794365: F external/org_tensorflow/tensorflow/core/tpu/tpu_library_init_fns.inc:34] TpuEmbeddingEngineState_Create not available in this library.

A bit more information on the TPU instance:

(py37) artem@t1v-n-bla-bla-bla-host:~/code/jax$ lsb_release -a
No LSB modules are available.
Distributor ID: Ubuntu
Description:    Ubuntu 20.04.2 LTS
Release:        20.04
Codename:       focal

JAX SHA: efa5edfd39632c5864e121de3290326b42b26326