Stuck on an issue?

Lightrun Answers was designed to reduce the constant googling that comes with debugging 3rd party libraries. It collects links to all the places you might be looking at while hunting down a tough bug.

And, if you’re still stuck at the end, we’re happy to hop on a call to see how we can help out.

Import error after compiling JAX from source on a TPU instance

See original GitHub issue

Discussed in https://github.com/google/jax/discussions/9219

^{Originally posted by awav January 17, 2022} 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-0:~/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</div>

Issue Analytics

State:
Created 2 years ago
Comments:11 (2 by maintainers)

Top GitHub Comments

3reactions

skyecommented, Jan 18, 2022

Hi @awav, try this with your custom-built jaxlib:

pip install --upgrade libtpu-nightly -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

Or that doesn’t work, pick the same libtpu-nightly date that you built your jaxlib:

pip install --upgrade libtpu-nightly==0.1.devYYYYMMDD -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

(Replace the YYYYMMDD with the build date)

Explanation: the TpuEmbeddingEngineState_Create not available in this library message means the jaxlib you compiled requires a newer version of libtpu, the library that provides low-level access to Cloud TPUs (sorry the error message is unclear!). By default, we automatically install compatible versions of jaxlib and libtpu-nightly when you pip install jax[tpu]. However, when you build your own jaxlib, you’ll have to manualy upgrade libtpu-nightly as shown above.

1reaction

awavcommented, Jan 20, 2022

Btw, I think this issue should be closed, the build works and I can import JAX w/o the issue. I opened another one #9261, and it is connected to the XLA and TPU interaction.

Top Results From Across the Web

Building from source - JAX documentation

Building JAX involves two steps: Building or installing jaxlib , the C++ support library for jax . Installing the jax Python package.

15. JAX - Python Programming for Economics and Finance

JAX provides data types, functions and a compiler for fast linear algebra operations ... The following import is standard, replacing import numpy as...

Cloud TPU VM user's guide

This guide describes how to set up a Google Cloud project for using Cloud TPU VMs. It describes the commands for using Cloud...

Newest 'tpu' Questions - Stack Overflow

I created a TPU VM on GCP. I logged in via ssh and want to install some software. But I get the following...

XLA: Optimizing Compiler for Machine Learning - TensorFlow

Explicit compilation API offers a fine-grained control for choosing which functions should be compiled. For example, the following ...