Error importing jax after building from source

[JuanGMendoza]

I have built jaxlib from source on a supercomputer cluster using linux Centos 7 Haswell.

I used gcc 8.3.0 because using anything newer gave me the error message " unsupported GNU version! gcc versions later than 9 are not supported!"

I also used an "old" commit of jax (0b4b0ba) because I need jaxlib 0.3.15

I installed it with GPU capabilities, using cuda-11.6.2 and cudnn-8.4.0.27

When trying to import jax I get the following error message

ImportError: /.conda/envs/sourceXena_8.3.0/lib/python3.10/site-packages/jaxlib/xla_extension.so: undefined symbol: cublasGetStatusString

I have not been able to find anything regarding this. Any suggestions?

[hawkinsp]

The issue is that the particular commit you chose was in a temporarily broken state, since fixed. You need to use a newer commit that fixes the bug (or a preceding commit before the bug was introduced). The relevant commit will change the WORKSPACE file, so just look at the history of that file to find the subsequent commits that matter.

Hope that helps!

[JuanGMendoza]

@hawkinsp So then if I want jaxlib 0.3.15 should I just change jax/jax/version.py?

I had changed to the last commit before the version change just in case any other files needed to be different in order to have this version of jaxlib

[hawkinsp]

Try checking out the jaxlib-v0.3.15 tag?

You should not need to change anything, you just need to check out the right version.

[JuanGMendoza]

@hawkinsp sorry for the ignorance, but what do you mean by "checking out" the tag? Is this an argument in the build command?

[hawkinsp]

In your JAX git checkout, run git checkout jaxlib-v0.3.15

[JuanGMendoza]

Hello again Peter, I followed your instructions and built it again and I am receiving a similar error now @hawkinsp

jax._src.traceback_util.UnfilteredStackTrace: jaxlib.xla_extension.XlaRuntimeError: INTERNAL: CustomCall failed: jaxlib/cuda/cuda_prng_kernels.cc:32: operation cudaGetLastError() failed: cudaGetErrorString symbol not found.