New jaxlib wheels cause `ptxas` error

Hi everyone! I’ve been using (and enjoying 😁) JAX on an university cluster for quite a while, without much problems until now.

Unfortunately since the new jaxlib wheels 0.1.66 and 0.1.67 have been released I’m not able to upgrade and keep a functioning install.

Here’s the error I get after upgrading

>>> import jax, jaxlib
>>> print(jax.version.__version__, jaxlib.version.__version__)
0.2.13 0.1.67
>>> a = jax.numpy.ones((100,100))**2
2021-05-20 10:09:06.617783: E external/org_tensorflow/tensorflow/stream_executor/gpu/asm_compiler.cc:93] You are using ptxas 8.x, but TF requires ptxas 9.x (and strongly prefers >= 9.2.88).  Compilation of XLA kernels below will likely fail.

You do not need to update CUDA; cherry-picking the ptxas binary is sufficient.
2021-05-20 10:09:06.630330: F external/org_tensorflow/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc:478] ptxas returned an error during compilation of ptx to sass: 'Internal: ptxas exited with non-zero error code 65280, output: ptxas /tmp/tempfile-kinsler-56b28778-3690660-5c2bf4b9c1e23, line 5; fatal   : Unsupported .version 6.0; current version is '5.0'
ptxas fatal   : Ptx assembly aborted due to errors
'  If the error message indicates that a file could not be written, please verify that sufficient filesystem space is provided.
Aborted (core dumped)

Reverting back to jaxlib 0.1.65 fixes the problem. The CUDA installation is non-standard (the cluster uses environment modules and I don’t have admin rights), but by setting the appropriate environment variables and paths I’ve got everything working until the new wheels came out. For completeness, here’s my nvcc

$ nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Thu_Jan_28_19:32:09_PST_2021
Cuda compilation tools, release 11.2, V11.2.142
Build cuda_11.2.r11.2/compiler.29558016_0

[hawkinsp]

#6820 should have fixed this issue, although you’ll either need to rebuild jaxlib from source or wait for us to make a new release. Hope that helps!

[astanziola]

@hawkinsp That is amazing!! Thank you very much for helping with this, as soon as the new release is out I’ll check it out (I can’t compile on the cluster)