BUG: ipu_xla_client_pybind.so: undefined symbol

[PierrickPochelu]

Description

I did
" pip install jax==0.3.16+ipu jaxlib==0.3.15+ipu.sdk310 -f https://graphcore-research.github.io/jax-experimental/wheels.html "
as mentioned in the documentation.

When I open a Python session and try to "import jax" I get the following error message:

[...]
    from . import ipu_xla_client_pybind as _ipu_xla
ImportError: /home/ipuuser/manu_tests/.venv_pytorch/lib/python3.8/site-packages/jaxlib/ipu_xla_client_pybind.so: undefined symbol: _ZNK6poplar8IPUModel12createDeviceENS_11OptionFlagsEbj, version V0

What jax/jaxlib version are you using?

JAX 0.3.16+ipu , JAXLIB .3.15+ipu.sdk310

Which accelerator(s) are you using?

IPU

Additional System Info

Python 3.8.10, Ubuntu 20.04.5 LTS

[balancap]

Sorry to hear it did not work out of the box!

Do you know which version of the Poplar SDK you are using? jaxlib is compiled against Poplar C++ libraries, so will need a local install of the later. In my experience, the easiest way is run inside one of the Graphcore docker image (e.g. https://hub.docker.com/layers/graphcore/pytorch/3.2.0-ubuntu-20.04/images/sha256-30f0e08f260b76f2821c66cbd7dbe78d4a2a85e040d10185a03a0dbd1499ff2e?context=explore)

[PierrickPochelu]

We have 3.0.0 poplar SDK on our Graphcore server. It is enough for our scientific works on PyTorch/Tensorflow.

If you guide us for having Cloud access to some IPUs we could try your last version of IPUs & poplar. It may be very useful for our PhD students, and our research, where we plan to use more and more Jax.

What do you think ?

Best regards,
Pierrick Pochelu

[balancap]

I could reproduce the issue with SDK 3.0. I'll give it a try at compiling jaxlib with it, but I can't guarantee a successful outcome.
In the meantime, if you want to try IPUs with latest SDK version, the easiest way for you (and students) is to use IPU on Paperspace. Once logged, you'll be able to run notebooks with IPUs. For instance, here is the quickstart JAX on IPU notebook: https://console.paperspace.com/github/graphcore-research/jax-experimental?container=graphcore%2Fpytorch-jupyter%3A3.1.0-ubuntu-20.04&machine=Free-IPU-POD4&file=%2Fipu%2Fexamples%2Fjax_ipu_quickstart.ipynb

Hopefully this one should run out of the box!

[PierrickPochelu]

I didn t succeeded to log in. The blue "sign in" refresh the page. I tried on Ubuntu & Windows: Chrome, Firefox, Edge.

[balancap]

Sorry to hear, I have forwarded the issue to the internal Paperspace team. It may help to first create & login on the main page (https://www.paperspace.com/), and then reload the notebook?

In the meantime, I have also manage to compile jaxlib for Poplar SDK 3.0. The following install should now work:

pip install jax==0.3.16+ipu jaxlib==0.3.15+ipu.sdk300 -f https://graphcore-research.github.io/jax-experimental/wheels.html

All basics should be working properly. Please tell us if there is any feature we can improve!

[PierrickPochelu]

I saw you have just compiled and published the 3.0 wheel, congratulations! Jax is working fine.

Regarding my cloud account, I have an issue with Paperspace account and not related to you. The graphcore-jax notebook on the cloud is working fine with my colleague account.

Thank you!

[balancap]

Thanks 👍 Please feel free to post any issue you discover, I'll try my best to help :)